Skip to main content

scouter_sql/sql/traits/
tag.rs

1use crate::sql::error::SqlError;
2use crate::sql::query::Queries;
3
4use async_trait::async_trait;
5use itertools::multiunzip;
6use scouter_types::{Tag, TagRecord};
7use sqlx::{postgres::PgQueryResult, types::Json, Pool, Postgres};
8use std::result::Result::Ok;
9use tracing::error;
10
11#[async_trait]
12pub trait TagSqlLogic {
13    /// Attempts to insert multiple tag records into the database in a batch.
14    ///
15    /// # Arguments
16    /// * `pool` - The database connection pool
17    /// * `baggage` - The trace baggage records to insert
18    async fn insert_tag_batch(
19        pool: &Pool<Postgres>,
20        tags: &[TagRecord],
21    ) -> Result<PgQueryResult, SqlError> {
22        let query = Queries::InsertTag.get_query();
23
24        let (entity_type, entity_id, key, value): (Vec<&str>, Vec<&str>, Vec<&str>, Vec<&str>) =
25            multiunzip(tags.iter().map(|b| {
26                (
27                    b.entity_type.as_str(),
28                    b.entity_id.as_str(),
29                    b.key.as_str(),
30                    b.value.as_str(),
31                )
32            }));
33
34        let query_result = sqlx::query(query)
35            .bind(entity_type)
36            .bind(entity_id)
37            .bind(key)
38            .bind(value)
39            .execute(pool)
40            .await
41            .inspect_err(|e| error!("Error inserting tags: {:?}", e))?;
42
43        Ok(query_result)
44    }
45
46    async fn get_tags(
47        pool: &Pool<Postgres>,
48        entity_type: &str,
49        entity_id: &str,
50    ) -> Result<Vec<TagRecord>, SqlError> {
51        let query = Queries::GetTags.get_query();
52
53        let rows = sqlx::query_as::<_, TagRecord>(query)
54            .bind(entity_type)
55            .bind(entity_id)
56            .fetch_all(pool)
57            .await?;
58
59        Ok(rows)
60    }
61
62    async fn get_entity_id_by_tags(
63        pool: &Pool<Postgres>,
64        entity_type: &str,
65        tags: &[Tag],
66        match_all: bool,
67    ) -> Result<Vec<String>, SqlError> {
68        let query = Queries::GetEntityIdByTags.get_query();
69
70        let rows = sqlx::query_scalar::<_, String>(query)
71            .bind(entity_type)
72            .bind(Json(tags))
73            .bind(match_all)
74            .fetch_all(pool)
75            .await?;
76
77        Ok(rows)
78    }
79}