kafru/
agent.rs

1use core::fmt;
2use chrono::{DateTime, Utc};
3use serde::{Deserialize, Serialize};
4use surrealdb::RecordId;
5use crate::database::Db;
6use std::sync::Arc;
7use tracing::{info, error, instrument};
8use std::sync::atomic::{AtomicBool, Ordering};
9use crate::Command;
10
11/// An enum representing the types of agents.
12///
13/// This enum is used to classify the different kinds of agents in the system. Each variant 
14/// represents a specific role or functionality that the agent serves.
15///
16#[derive(Debug,Clone, PartialEq, Deserialize, Serialize)]
17pub enum  AgentKind {
18    Queue,
19    Scheduler,
20    Task,
21    None
22}
23
24impl fmt::Display for AgentKind {
25
26    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        
27        match self {
28            Self::Scheduler => write!(f,"Scheduler"),
29            Self::Queue => write!(f,"Queue"),
30            Self::Task => write!(f,"Task"),
31            Self::None => write!(f,"None")
32        }
33    }
34}
35
36/// An enum representing the possible statuses of an agent.
37///
38/// This enum defines the different states that an agent can be in, providing a way to track 
39/// the lifecycle or status of an agent during its execution. The statuses can be used to indicate 
40/// whether the agent has been initialized, is actively running, has been terminated, encountered an 
41/// error, or has completed its task.
42///
43#[derive(Debug,Clone, PartialEq, Deserialize, Serialize)]
44pub enum AgentStatus {
45    Initialized,
46    Running,
47    Terminated,
48    Error,
49    Completed
50}
51
52impl fmt::Display for AgentStatus {
53    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
54        match self {
55            Self::Initialized => write!(f,"Initialized"),
56            Self::Running => write!(f,"Running"),
57            Self::Terminated => write!(f,"Terminated"),
58            Self::Error => write!(f,"Error"),
59            Self::Completed => write!(f,"Completed"),
60        }
61    }
62}
63
64/// A struct representing the data of an agent.
65///
66/// This struct holds the detailed information of an agent, including optional fields such as its 
67/// unique identifier, task ID, status, commands, and server-related information. The fields are 
68/// serializable using `serde`, and the `skip_serializing_if` attribute ensures that fields with 
69/// `None` values are not included during serialization.
70///
71#[derive(Debug,Clone,Deserialize,Serialize,PartialEq)]
72pub struct AgentData {
73    #[serde(skip_serializing_if="Option::is_none")]
74    pub id: Option<RecordId>,
75    #[serde(skip_serializing_if="Option::is_none")]
76    pub parent_id: Option<RecordId>,
77    #[serde(skip_serializing_if="Option::is_none")]
78    pub queue_id: Option<RecordId>,
79    #[serde(skip_serializing_if="Option::is_none")]
80    pub name: Option<String>,
81    #[serde(skip_serializing_if="Option::is_none")]
82    pub kind: Option<AgentKind>,
83    #[serde(skip_serializing_if="Option::is_none")]
84    pub server: Option<String>,
85    #[serde(skip_serializing_if="Option::is_none")]
86    pub task_id: Option<u64>,
87    #[serde(skip_serializing_if="Option::is_none")]
88    pub status: Option<AgentStatus>,
89    #[serde(skip_serializing_if="Option::is_none")]
90    pub command: Option<Command>,
91    #[serde(skip_serializing_if="Option::is_none")]
92    pub command_is_executed: Option<bool>,
93    #[serde(skip_serializing_if="Option::is_none")]
94    pub message: Option<String>,
95    #[serde(skip_serializing_if="Option::is_none")]
96    pub author: Option<String>,
97    pub date_modified: DateTime<Utc>
98}
99
100impl Default for AgentData {
101    fn default() -> Self {
102        Self {
103            id: None,
104            parent_id: None,
105            queue_id: None,
106            name: None,
107            kind: None,
108            server: None,
109            task_id: None,
110            status: None,
111            command: None,
112            command_is_executed: None,
113            message: None,
114            author: None,
115            date_modified: Utc::now()
116        }
117    }
118}
119
120/// A struct representing the filters that can be applied when querying agent records.
121///
122/// This struct allows specifying various optional fields that can be used to filter the agents 
123/// based on different attributes such as their ID, name, status, task ID, and more. Each field 
124/// is optional, meaning the filter can be customized based on the requirements of the query.
125///
126#[derive(Debug,Clone,Deserialize,Serialize, Default)]
127pub struct AgentFilter {
128    pub id: Option<RecordId>,
129    pub ids: Option<Vec<RecordId>>,
130    pub parent_id: Option<RecordId>,
131    pub parent_ids: Option<Vec<RecordId>>,
132    pub queue_id: Option<RecordId>,
133    pub queue_ids: Option<Vec<RecordId>>,
134    pub name: Option<String>,
135    pub names: Option<Vec<String>>,
136    pub kind: Option<AgentKind>,
137    pub server: Option<String>,
138    pub task_id: Option<u64>,
139    pub status: Option<AgentStatus>,
140    pub statuses: Option<Vec<AgentStatus>>,
141    pub command: Option<Command>,
142    pub commands: Option<Vec<Command>>,
143    pub command_is_executed: Option<bool>,
144    pub message: Option<String>,
145}
146
147#[derive(Debug, Clone)]
148pub struct Agent {
149    db: Arc<Db>,
150    table: String,
151}
152
153static AGENT_TABLE_CREATED: AtomicBool = AtomicBool::new(false);
154
155impl Agent {
156
157    /// Creates a new instance of the struct with an optional database connection.
158    ///
159    /// This function initializes a new instance of the struct by either using the provided database connection (`db`) 
160    /// or by creating a new one if `db` is `None`. The database connection is wrapped in an `Arc` for shared ownership.
161    /// It also sets a default table name (`"kafru_agents"`) for database operations.
162    ///
163    /// # Arguments
164    /// * `db` - An optional `Arc<Db>` representing the database connection. If `None` is provided, a new database connection 
165    ///   will be created.
166    ///
167    /// # Returns
168    /// * `Self` - A new instance of the struct initialized with the provided or newly created database connection and 
169    ///   the default table name.
170    ///
171    pub async fn new(db: Option<Arc<Db>>) -> Self {
172        let db: Arc<Db> = if let Some(value) = db {
173            value.clone()
174        }
175        else {
176            Arc::new(Db::new(None).await.unwrap())
177        };
178        let table: String = "kafru_agents".to_string();
179        Self {
180            db,
181            table
182        }
183    }
184
185    /// Purges agent data for a specific server.
186    ///
187    /// This function deletes all agent records associated with a given server from the database. The operation is performed
188    /// on the table specified by `self.table`, which is dynamically bound to the query. If the operation is successful, the
189    /// function returns `Ok(true)`. If an error occurs during the database query, an error message is returned.
190    ///
191    /// # Arguments
192    /// * `server` - The name of the server for which agent data will be purged. All records associated with this server will be deleted.
193    ///
194    /// # Returns
195    /// * `Ok(true)` - If the data was successfully purged for the given server.
196    /// * `Err(String)` - An error message if the purge operation fails.
197    ///
198    #[instrument(skip_all)]
199    pub async fn purge(&self,server: String) -> Result<bool,String> {
200        match self.db.client.query("DELETE type::table($table) WHERE server = $server")
201            .bind(("table",self.table.clone()))
202            .bind(("server",server.clone()))
203        .await {
204            Ok(_) => {
205                Ok(true)
206            },
207            Err(error) => {
208                error!("{}",error);
209                Err(format!("unable to purge {} for server {}",self.table,server))
210            }
211        }
212    }
213    
214    /// Creates a table for agent data if it does not already exist.
215    ///
216    /// This function attempts to create a table in the database using the name stored in `self.table`. If the table
217    /// has already been created, it simply returns `true`. If the table creation is successful, it proceeds to purge
218    /// old agent data and updates a global flag `AGENT_TABLE_CREATED` to indicate that the table is now created.
219    ///
220    /// # Arguments
221    /// * `server` - The name of the server for which the table should be created. This is used later in the function
222    ///   to trigger any additional purging of related data.
223    ///
224    /// # Returns
225    /// * `Ok(true)` - If the table was successfully created or already exists, and the data purging operation was successful.
226    /// * `Err(String)` - An error message if the table creation fails or any related data purging fails.
227    ///
228    #[instrument(skip_all)]
229    pub async fn create_table(&self, server: String) -> Result<bool,String> {
230        let is_created: bool = AGENT_TABLE_CREATED.load(Ordering::Relaxed);
231        if is_created {
232            return Ok(true);
233        }
234        let stmt: String = format!("DEFINE TABLE IF NOT EXISTS {};",self.table);
235        match self.db.client.query(stmt).await {
236            Ok(_) => {
237                info!("table {} has been created",&self.table);
238                if let Err(error) = self.purge(server).await {                
239                    return Err(error);
240                }
241                AGENT_TABLE_CREATED.store(true, Ordering::Relaxed);
242                Ok(true)
243            }
244            Err(error) => {
245                error!("{}",error);
246                Err(format!("unable to create table {}",&self.table))
247            }
248        }
249    }
250
251    /// Removes an agent record by its unique identifier, with an option to delete related records.
252    ///
253    /// This function deletes an agent record from the database using the provided `id`. If the `include_related`
254    /// flag is set to `true`, it will also delete related records that reference this agent by its `parent_id`.
255    /// If successful, it returns `true`, otherwise an error message is returned.
256    ///
257    /// # Arguments
258    /// * `id` - The unique identifier (`RecordId`) of the agent to remove.
259    /// * `include_related` - A boolean flag indicating whether related records should also be deleted.
260    ///   - If `true`, it will delete any records with a `parent_id` matching the agent's `id`.
261    ///   - If `false`, only the agent record will be deleted.
262    ///
263    /// # Returns
264    /// * `Ok(true)` - If the agent record (and optionally related records) was successfully deleted.
265    /// * `Err(String)` - An error message if the deletion fails.
266    ///
267    pub async fn remove(&self,id: RecordId, include_related: bool) -> Result<bool,String> {
268        match self.db.client.delete::<Option<AgentData>>(id.clone()).await {
269            Ok(_) => {
270                if include_related {
271                    if let Err(error)  = self.db.client.query("DELETE FROM type::table($table) WHERE parent_id=type::thing($parent_id)")
272                    .bind(("table",self.table.clone()))
273                    .bind(("parent_id",id.clone())).await {
274                        error!("{}",error);
275                        return Err(format!("database error when deleting related data of agent data with id: {:?}",id));    
276                    }
277                }
278                Ok(true)
279            }
280            Err(error) => {
281                error!("{}",error);
282                Err(format!("database error when deleting agent data with id: {:?}",id))
283            }
284        }
285    }
286    
287    /// Retrieves an agent record by its unique identifier.
288    ///
289    /// This function queries the database for an agent record using the provided `id`. If the record
290    /// exists, it is returned. Otherwise, an error is returned indicating that the agent could not be found.
291    ///
292    /// # Arguments
293    /// * `id` - The unique identifier (`RecordId`) of the agent to retrieve.
294    ///
295    /// # Returns
296    /// * `Ok(AgentData)` - The agent record if found.
297    /// * `Err(String)` - An error message if the agent record could not be found or if a database error occurs.
298    ///
299    pub async fn get(&self,id: RecordId) -> Result<AgentData, String> {
300
301        match self.db.client.select::<Option<AgentData>>(id.clone()).await {
302            Ok(data) => {
303                if let Some(item) = data {
304                    return Ok(item);
305                }
306                return Err(format!("unable to find agent with id: {:?}",id));
307            }
308            Err(error) => {
309                error!("{}",error);
310                Err(format!("database error when retrieving agent data with id: {:?}",id))
311            }
312        }
313    }
314
315    /// Retrieves an agent record by its name and associated server.
316    ///
317    /// This function queries the database to fetch an agent record based on the provided `name`
318    /// and `server` values. It uses the specified table in the database for the query.
319    ///
320    /// # Arguments
321    /// * `name` - The name of the agent to search for.
322    /// * `server` - The server associated with the agent.
323    ///
324    /// # Returns
325    /// * `Ok(AgentData)` if a matching agent record is found.
326    /// * `Err(String)` if no matching record is found or if a database error occurs.
327    ///
328    pub async fn get_by_name(&self,name: String,server: String) -> Result<AgentData, String> {
329        match self.db.client.query("SELECT * FROM type::table($table) WHERE name=$name AND server=$server")
330        .bind(("table",self.table.clone()))
331        .bind(("name",name.clone()))
332        .bind(("server",server.clone())).await {
333            Ok(mut response) => {
334                if let Ok(data) = response.take::<Option<AgentData>>(0) {
335                    if let Some(item) = data {
336                        return Ok(item);
337                    }
338                }
339                return Err(format!("unable to find agent with name: {}, server: {}",name, server));
340            }
341            Err(error) => {
342                error!("{}",error);
343                Err(format!("database error when retrieving agent data with name:{}, server: {}",name, server))
344            }
345        }
346    }
347
348    /// Lists agent records based on the provided filters.
349    ///
350    /// This function queries the database to retrieve a list of agents that match the specified
351    /// criteria. The filters are used to dynamically construct the query with optional conditions.
352    ///
353    /// # Arguments
354    /// * `filters` - An instance of `AgentFilter`
355    ///
356    /// # Returns
357    /// * `Ok(Vec<AgentData>)` - A vector of `AgentData` objects matching the filters.
358    /// * `Err(String)` - An error message if the query fails or if data retrieval encounters an issue.
359    ///
360    pub async fn list(&self,filters: AgentFilter) -> Result<Vec<AgentData>,String> {
361        let mut  stmt: String = "SELECT * FROM  type::table($table)".to_string();
362        let mut where_stmt: Vec<String> = Vec::new();
363
364        // WHERE placeholders
365        if filters.id.is_some() {
366            where_stmt.push("type::thing(id)=$id".to_string());
367        }
368        if filters.parent_id.is_some() {
369            where_stmt.push("parent_id=type::thing($parent_id)".to_string());
370        }
371        if filters.command.is_some() {
372            where_stmt.push("command=$command".to_string());
373        }
374        if filters.command_is_executed.is_some() {
375            where_stmt.push("command_is_executed=$command_is_executed".to_string());
376        }
377        if filters.kind.is_some() {
378            where_stmt.push("kind=$kind".to_string());
379        }
380        if filters.name.is_some() {
381            where_stmt.push("name=$name".to_string());
382        }
383        if filters.queue_id.is_some() {
384            where_stmt.push("queue_id=$queue_id".to_string());
385        }
386        if filters.server.is_some() {
387            where_stmt.push("server=$server".to_string());
388        }
389        if filters.status.is_some() {
390            where_stmt.push("status=$status".to_string());
391        }
392        if filters.statuses.is_some() {
393            where_stmt.push("status IN $statuses".to_string());
394        }
395        if filters.commands.is_some() {
396            where_stmt.push("command IN $commands".to_string());
397        }
398        if filters.names.is_some() {
399            where_stmt.push("name IN $names".to_string());
400        }
401        if filters.queue_ids.is_some() {
402            where_stmt.push("queue_id IN $queue_ids".to_string());
403        }
404        if filters.ids.is_some() {
405            where_stmt.push("id IN $ids".to_string());
406        }
407        if !where_stmt.is_empty() {
408            stmt = format!("{} WHERE {}",stmt,where_stmt.join(" AND "));
409        }
410
411        // VALUE Binding
412        let mut query = self.db.client.query(stmt).bind(("table",self.table.clone()));
413        
414        if let Some(value) = filters.id {
415            query = query.bind(("id",value));
416        }
417        if let Some(value) = filters.parent_id {
418            query = query.bind(("parent_id",value));
419        }
420        if let Some(value) = filters.command {
421            query = query.bind(("command",value));
422        }
423        if let Some(value) = filters.command_is_executed {
424            query = query.bind(("command_is_executed",value));
425        }
426        if let Some(value) = filters.kind {
427            query = query.bind(("kind",value));
428        }
429        if let Some(value) = filters.name {
430            query = query.bind(("name",value));
431        }
432        if let Some(value) = filters.queue_id {
433            query = query.bind(("queue_id",value));
434        }
435        if let Some(value) = filters.server {
436            query = query.bind(("server",value));
437        }
438        if let Some(value) = filters.status {
439            query = query.bind(("status",value));
440        }
441        if let Some(values) = filters.statuses {
442            query = query.bind(("statuses",values));
443        }
444        if let Some(values) = filters.commands {
445            query = query.bind(("commands",values));
446        }
447        if let Some(values) = filters.queue_ids {
448            query = query.bind(("queue_ids",values));
449        }
450        if let Some(values) = filters.ids {
451            query = query.bind(("ids",values));
452        }
453        if let Some(values) = filters.names {
454            query = query.bind(("names",values));
455        }
456        match query.await {
457            Ok(mut response) => {
458                if let Ok(data) =  response.take::<Vec<AgentData>>(0) {
459                    return Ok(data);
460                }
461                return Err(format!("unable to retrive agent data"));
462            }
463            Err(error) => {
464                error!("{}",error);
465                return Err("database error when retrieving agent data".to_string());
466            }
467        }
468    }
469
470    /// Updates an agent record in the database.
471    ///
472    /// This function retrieves the current record associated with the given `id` and updates its fields
473    /// with the values provided in `data`. Any fields in `data` that are `None` will retain the values
474    /// from the existing record.
475    ///
476    /// # Arguments
477    /// * `id` - The unique identifier (`RecordId`) of the agent record to update.
478    /// * `data` - An `AgentData` instance containing the updated field values. Fields with `None`
479    ///   values will not overwrite the existing data.
480    ///
481    /// # Returns
482    /// * `Ok(AgentData)` - The updated `AgentData` record after successfully applying the changes.
483    /// * `Err(String)` - An error message if the update fails or if the record is not found.
484    ///
485    pub async fn update(&self,id: RecordId, data:AgentData) -> Result<AgentData,String> {
486        match self.get(id.clone()).await  {
487            Ok(record)=> {
488                let data: AgentData = AgentData {                            
489                    name: if data.name.is_none() { record.name } else { data.name },
490                    server: if data.server.is_none() { record.server } else {data.server },
491                    parent_id: if data.parent_id.is_none() { record.parent_id } else {data.parent_id },
492                    kind: if data.kind.is_none() { record.kind } else {data.kind },
493                    queue_id: if data.queue_id.is_none() { record.queue_id } else {data.queue_id },
494                    status: if data.status.is_none() { record.status } else {data.status },
495                    message: if data.message.is_none() { record.message } else {data.message },
496                    author: if data.author.is_none() { record.author } else {data.author },
497                    command: if data.command.is_none() { record.command } else {data.command },
498                    command_is_executed: if data.command_is_executed.is_none() { record.command_is_executed } else {data.command_is_executed },
499                    task_id: if data.task_id.is_none() { record.task_id } else {data.task_id },
500                    date_modified: Utc::now(),
501                    ..Default::default()
502                };
503                match self.db.client.update::<Option<AgentData>>(id).content(data.clone()).await {
504                    Ok(response) => {
505                        if let Some(data) = response {
506                            return Ok(data);
507                        }
508                        Err(format!("agent {} under server {} not found",data.name.unwrap(),data.server.unwrap()))
509                    }
510                    Err(error) => {
511                        error!("{}",error);
512                        Err(format!("unable to update agent {} under server {}",data.name.unwrap(),data.server.unwrap()))
513                    }
514                }
515            }
516            Err(error) => Err(error.to_string())
517        }
518    }
519
520    /// Sends a command to an agent by updating its record in the database.
521    ///
522    /// This function sets a command on an agent record, along with an optional message and author.
523    /// The `command_is_executed` field is reset to `false` to indicate the command needs to be executed.
524    ///
525    /// # Arguments
526    /// * `id` - The unique identifier (`RecordId`) of the agent to update.
527    /// * `command` - The command to be assigned to the agent.
528    /// * `message` - An optional message providing context or details about the command.
529    /// * `author` - An optional author name or identifier for tracking who issued the command.
530    ///
531    /// # Returns
532    /// * `Ok(AgentData)` - The updated agent record reflecting the assigned command.
533    /// * `Err(String)` - An error message if the update fails or if the record is not found.
534    ///
535    pub async fn send_command(&self,id: RecordId, command:Command, message: Option<String>, author: Option<String>) -> Result<AgentData,String> {
536
537        let data: AgentData = AgentData {                            
538            command: Some(command),
539            command_is_executed: Some(false),
540            message,
541            author,
542            ..Default::default()
543        };
544        match self.db.client.update::<Option<AgentData>>(id).merge(data.clone()).await {
545            Ok(response) => {
546                if let Some(data) = response {
547                    return Ok(data);
548                }
549                Err(format!("agent record not found"))
550            }
551            Err(error) => {
552                error!("{}",error);
553                Err(format!("unable to send command to agent"))
554            }
555        }
556    }
557    /// Registers a new agent in the database or updates an existing one if it already exists.
558    ///
559    /// This function performs validation on the provided `AgentData`, ensures the necessary database
560    /// table exists, and creates a new agent record. If an agent with the same `name` and `server` 
561    /// already exists, the existing record is removed before creating the new one.
562    ///
563    /// # Arguments
564    /// * `data` - An `AgentData` instance containing the details of the agent to register.
565    ///
566    /// # Returns
567    /// * `Ok(AgentData)` - The newly created or updated agent record.
568    /// * `Err(String)` - An error message if registration fails or the input data is invalid.
569    ///
570    pub async fn register(&self, mut data: AgentData) -> Result<AgentData,String>{
571        if let Some(value) = data.name.clone() {
572            if value.len() < 3 {
573                return Err("name must be atleast more than 3 characters".to_string());
574            }
575        }
576        else {            
577            return Err("name is required".to_string());
578        }
579        if data.kind.is_none() {
580            return Err("kind is required".to_string())
581        }
582        if data.task_id == Some(0) && data.kind == Some(AgentKind::Task) {
583            return Err(format!("task_id must be greater than 0 for task got {}",data.task_id.unwrap()));
584        }
585        if data.server.is_none() {
586            return Err(format!("server is required"));
587        }
588        if let Err(error) = self.create_table(data.server.clone().unwrap()).await {
589            return Err(error);
590        }
591        if data.command_is_executed.is_none() {
592            data.command_is_executed = Some(false);
593        }
594        if let Ok(items) = self.list(AgentFilter {
595            server: data.server.clone(),
596            name: data.name.clone(),
597            ..Default::default()
598        }).await {
599            if items.len() > 0 {
600                if let Some(item) = items.first() {
601                    self.remove(item.id.clone().unwrap(),true).await?;
602                }
603            }
604        }
605        match self.db.client.create::<Option<AgentData>>(self.table.clone()).content(data.clone()).await {
606            Ok(response) => {
607                if let Some(data) = response {
608                    return Ok(data);
609                }
610                return Err(format!("no data found for agent {} at server {}",data.name.unwrap(),data.server.unwrap()))
611            }
612            Err(error) => {
613                error!("{}",error);
614                Err(format!("unable to register agent {} at server {}",data.name.unwrap(),data.server.unwrap()))
615            }
616        }
617    }
618
619    
620    /// Converts a Tokio task ID into a `u64` value.
621    ///
622    /// # Arguments
623    /// * `id` - The Tokio task ID to be converted.
624    ///
625    /// # Returns
626    /// A `u64` representation of the given task ID.
627    pub async fn to_id(id: tokio::task::Id) -> u64 {
628        id.to_string().parse::<u64>().unwrap()
629    }
630
631    /// Generates a unique name for a task.
632    ///
633    /// # Arguments
634    /// * `queue_name` - A reference to the name of the queue the task belongs to.
635    /// * `task_id` - A unique identifier for the task.
636    ///
637    /// # Returns
638    /// A string combining the queue name and task ID in the format `{queue_name}-{task_id}`.
639    pub async fn to_name(queue_name: &String, task_id: &u64) -> String {
640        format!("{}-{}", queue_name, task_id)
641    }
642}
643
kafru/agent.rs

kafru/
agent.rs
