scouter_sql/sql/traits/
entity.rs

1use crate::sql::error::SqlError;
2use crate::sql::query::Queries;
3use async_trait::async_trait;
4use sqlx::{Pool, Postgres, Row};
5use std::result::Result::Ok;
6use tracing::{error, instrument};
7
8#[async_trait]
9pub trait EntitySqlLogic {
10    /// Get entity ID from UID
11    /// # Arguments
12    /// * `uid` - The UID of the entity
13    /// # Returns
14    /// * `Result<i32, SqlError>` - Result of the query returning the entity ID
15    #[instrument(skip_all)]
16    async fn get_entity_id_from_uid(pool: &Pool<Postgres>, uid: &str) -> Result<i32, SqlError> {
17        let query = Queries::GetEntityIdFromUid.get_query();
18
19        let result = sqlx::query(query)
20            .bind(uid)
21            .fetch_one(pool)
22            .await
23            .inspect_err(|e| error!("Failed to fetch entity ID from UID: {:?} uid: {}", e, uid))?;
24
25        let id: i32 = result.get("id");
26
27        Ok(id)
28    }
29
30    async fn get_optional_entity_id_from_uid(
31        pool: &Pool<Postgres>,
32        uid: &str,
33    ) -> Result<Option<i32>, SqlError> {
34        let query = Queries::GetEntityIdFromUid.get_query();
35
36        sqlx::query(query)
37            .bind(uid)
38            .fetch_optional(pool)
39            .await
40            .map_err(SqlError::SqlxError)
41            .map(|row| row.map(|r| r.get("id")))
42    }
43
44    async fn get_entity_id_from_space_name_version_drift_type(
45        pool: &Pool<Postgres>,
46        space: &str,
47        name: &str,
48        version: &str,
49        drift_type: &str,
50    ) -> Result<i32, SqlError> {
51        let query = Queries::GetEntityIdFromSpaceNameVersionDriftType.get_query();
52
53        let result = sqlx::query(query)
54            .bind(space)
55            .bind(name)
56            .bind(version)
57            .bind(drift_type)
58            .fetch_one(pool)
59            .await
60            .map_err(SqlError::SqlxError)?;
61
62        let id: i32 = result.get("id");
63
64        Ok(id)
65    }
66
67    /// Helper function to create a new entity
68    /// # Arguments
69    /// * `space` - The space of the entity
70    /// * `name` - The name of the entity
71    /// * `version` - The version of the entity
72    /// * `drift_type` - The drift type of the entity
73    /// # Returns
74    /// * `Result<String, SqlError>` - Result of the insert returning the new
75    async fn create_entity(
76        pool: &Pool<Postgres>,
77        space: &str,
78        name: &str,
79        version: &str,
80        drift_type: &str,
81    ) -> Result<(String, i32), SqlError> {
82        let query = "INSERT INTO scouter.drift_entities (space, name, version, drift_type) VALUES ($1, $2, $3, $4) ON CONFLICT (space, name, version, drift_type) DO NOTHING RETURNING id, uid;";
83
84        let result = sqlx::query(query)
85            .bind(space)
86            .bind(name)
87            .bind(version)
88            .bind(drift_type)
89            .fetch_one(pool)
90            .await
91            .map_err(SqlError::SqlxError)?;
92
93        let uid: String = result.get("uid");
94        let id: i32 = result.get("id");
95
96        Ok((uid, id))
97    }
98
99    async fn get_uid_from_args(
100        pool: &Pool<Postgres>,
101        space: &str,
102        name: &str,
103        version: &str,
104        drift_type: &str,
105    ) -> Result<String, SqlError> {
106        let query = "SELECT uid FROM entities WHERE space = $1 AND name = $2 AND version = $3 AND drift_type = $4;";
107
108        let result = sqlx::query(query)
109            .bind(space)
110            .bind(name)
111            .bind(version)
112            .bind(drift_type)
113            .fetch_one(pool)
114            .await
115            .map_err(SqlError::SqlxError)?;
116
117        let uid: String = result.get("uid");
118
119        Ok(uid)
120    }
121}
scouter_sql/sql/traits/entity.rs

scouter_sql/sql/traits/
entity.rs
