scouter_sql/sql/traits/
tag.rs

1use crate::sql::error::SqlError;
2use crate::sql::query::Queries;
3
4use async_trait::async_trait;
5use chrono::{DateTime, Utc};
6use itertools::multiunzip;
7use scouter_types::TagRecord;
8use sqlx::{postgres::PgQueryResult, Pool, Postgres};
9use std::result::Result::Ok;
10
11#[async_trait]
12pub trait TagSqlLogic {
13    /// Attempts to insert multiple trace baggage 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 (created_at, entity_type, entity_id, key, value): (
25            Vec<DateTime<Utc>>,
26            Vec<&str>,
27            Vec<&str>,
28            Vec<&str>,
29            Vec<&str>,
30        ) = multiunzip(tags.iter().map(|b| {
31            (
32                b.created_at,
33                b.entity_type.as_str(),
34                b.entity_id.as_str(),
35                b.key.as_str(),
36                b.value.as_str(),
37            )
38        }));
39
40        let query_result = sqlx::query(&query.sql)
41            .bind(created_at)
42            .bind(entity_type)
43            .bind(entity_id)
44            .bind(key)
45            .bind(value)
46            .execute(pool)
47            .await?;
48
49        Ok(query_result)
50    }
51
52    async fn get_tags(
53        pool: &Pool<Postgres>,
54        entity_type: &str,
55        entity_id: &str,
56    ) -> Result<Vec<TagRecord>, SqlError> {
57        let query = Queries::GetTags.get_query();
58
59        let rows = sqlx::query_as::<_, TagRecord>(&query.sql)
60            .bind(entity_type)
61            .bind(entity_id)
62            .fetch_all(pool)
63            .await?;
64
65        Ok(rows)
66    }
67}