boolean.expanded.rs

use serde::Deserialize as _;
#[pco_store::store(group_by = [database_id], float_round = 2)]
pub struct QueryStat {
    pub database_id: i64,
    pub toplevel: bool,
    pub calls: i64,
}
#[doc = concat!(
    " Generated by pco_store to store and load compressed versions of [",
    stringify!(QueryStat), "]"
)]
pub struct CompressedQueryStats {
    filter: Option<Filter>,
    database_id: i64,
    toplevel: Vec<u8>,
    calls: Vec<u8>,
}
impl CompressedQueryStats {
    pub fn new(rows: &Vec<QueryStat>) -> anyhow::Result<Self> {
        Ok(Self {
            database_id: rows[0].database_id.clone(),
            toplevel: ::pco::standalone::simple_compress(
                &rows.iter().map(|r| r.toplevel as u16).collect::<Vec<_>>(),
                &::pco::ChunkConfig::default(),
            )?,
            calls: ::pco::standalone::simple_compress(
                &rows.iter().map(|r| r.calls).collect::<Vec<_>>(),
                &::pco::ChunkConfig::default(),
            )?,
            filter: None,
        })
    }
    /// Loads data for the specified filters.
    pub async fn load(
        db: &impl ::std::ops::Deref<Target = deadpool_postgres::ClientWrapper>,
        mut filter: Filter,
        fields: impl TryInto<Fields>,
    ) -> anyhow::Result<Vec<CompressedQueryStats>> {
        let mut fields = fields
            .try_into()
            .map_err(|_| anyhow::Error::msg("unknown field"))?;
        fields.merge_filter(&filter);
        if filter.database_id.is_empty() {
            return Err(anyhow::Error::msg("database_id".to_string() + " is required"));
        }
        let sql = "SELECT ".to_string() + fields.select().as_str() + " FROM "
            + "query_stats" + " WHERE " + "database_id = ANY($1)";
        let mut results = Vec::new();
        for row in db
            .query(&db.prepare_cached(&sql).await?, &[&filter.database_id])
            .await?
        {
            results.push(fields.load_from_row(row, Some(filter.clone()))?);
        }
        Ok(results)
    }
    /// Deletes data for the specified filters, returning it to the caller.
    ///
    /// Note that all rows are returned from [decompress][Self::decompress] even if post-decompress filters would normally apply.
    pub async fn delete(
        db: &impl ::std::ops::Deref<Target = deadpool_postgres::ClientWrapper>,
        mut filter: Filter,
        fields: impl TryInto<Fields>,
    ) -> anyhow::Result<Vec<CompressedQueryStats>> {
        let mut fields = fields
            .try_into()
            .map_err(|_| anyhow::Error::msg("unknown field"))?;
        fields.merge_filter(&filter);
        if filter.database_id.is_empty() {
            return Err(anyhow::Error::msg("database_id".to_string() + " is required"));
        }
        let sql = "DELETE FROM ".to_string() + "query_stats" + " WHERE "
            + "database_id = ANY($1)" + " RETURNING " + fields.select().as_str();
        let mut results = Vec::new();
        for row in db
            .query(&db.prepare_cached(&sql).await?, &[&filter.database_id])
            .await?
        {
            results.push(fields.load_from_row(row, None)?);
        }
        Ok(results)
    }
    /// Decompresses a group of data points.
    pub fn decompress(self) -> anyhow::Result<Vec<QueryStat>> {
        let mut results = Vec::new();
        let toplevel: Vec<u16> = if self.toplevel.is_empty() {
            Vec::new()
        } else {
            ::pco::standalone::simple_decompress(&self.toplevel)?
        };
        let calls: Vec<i64> = if self.calls.is_empty() {
            Vec::new()
        } else {
            ::pco::standalone::simple_decompress(&self.calls)?
        };
        let len = [toplevel.len(), calls.len()].into_iter().max().unwrap_or(0);
        for index in 0..len {
            let row = QueryStat {
                database_id: self.database_id.clone(),
                toplevel: toplevel.get(index).cloned().unwrap_or_default() == 1,
                calls: calls.get(index).cloned().unwrap_or_default(),
            };
            if self.filter.as_ref().map(|f| f.matches(&row)) != Some(false) {
                results.push(row);
            }
        }
        Ok(results)
    }
    /// Writes the data to disk.
    pub async fn store(
        db: &impl ::std::ops::Deref<Target = deadpool_postgres::ClientWrapper>,
        rows: Vec<QueryStat>,
    ) -> anyhow::Result<()> {
        if rows.is_empty() {
            return Ok(());
        }
        let mut grouped_rows: ahash::AHashMap<_, Vec<QueryStat>> = ahash::AHashMap::new();
        for row in rows {
            grouped_rows.entry((row.database_id.clone(),)).or_default().push(row);
        }
        let sql = "COPY query_stats (database_id, toplevel, calls) FROM STDIN BINARY";
        let types = &[
            tokio_postgres::types::Type::INT8,
            tokio_postgres::types::Type::BYTEA,
            tokio_postgres::types::Type::BYTEA,
        ];
        let stmt = db.copy_in(&db.prepare_cached(&sql).await?).await?;
        let writer = tokio_postgres::binary_copy::BinaryCopyInWriter::new(stmt, types);
        futures::pin_mut!(writer);
        for rows in grouped_rows.into_values() {
            let row = Self::new(&rows)?;
            writer
                .as_mut()
                .write(&[&rows[0].database_id, &row.toplevel, &row.calls])
                .await?;
        }
        writer.finish().await?;
        Ok(())
    }
    /// Writes the data to disk, with the provided grouping closure applied.
    ///
    /// This can be used to improve the compression ratio and reduce read IO, for example
    /// by compacting real-time data into a single row per hour / day / week.
    pub async fn store_grouped<F, R>(
        db: &impl ::std::ops::Deref<Target = deadpool_postgres::ClientWrapper>,
        rows: Vec<QueryStat>,
        grouping: F,
    ) -> anyhow::Result<()>
    where
        F: Fn(&QueryStat) -> R,
        R: Eq + std::hash::Hash,
    {
        if rows.is_empty() {
            return Ok(());
        }
        let mut grouped_rows: ahash::AHashMap<_, Vec<QueryStat>> = ahash::AHashMap::new();
        for row in rows {
            grouped_rows
                .entry((row.database_id.clone(), grouping(&row)))
                .or_default()
                .push(row);
        }
        let sql = "COPY query_stats (database_id, toplevel, calls) FROM STDIN BINARY";
        let types = &[
            tokio_postgres::types::Type::INT8,
            tokio_postgres::types::Type::BYTEA,
            tokio_postgres::types::Type::BYTEA,
        ];
        let stmt = db.copy_in(&db.prepare_cached(&sql).await?).await?;
        let writer = tokio_postgres::binary_copy::BinaryCopyInWriter::new(stmt, types);
        futures::pin_mut!(writer);
        for rows in grouped_rows.into_values() {
            let row = Self::new(&rows)?;
            writer
                .as_mut()
                .write(&[&rows[0].database_id, &row.toplevel, &row.calls])
                .await?;
        }
        writer.finish().await?;
        Ok(())
    }
}
#[serde_with::serde_as]
#[derive(Debug, Default, serde::Deserialize, Clone, PartialEq)]
#[serde(deny_unknown_fields)]
#[doc = concat!(
    " Generated by pco_store to specify filters when loading [", stringify!(QueryStat),
    "]"
)]
pub struct Filter {
    #[serde(default)]
    #[serde_as(deserialize_as = "serde_with::DefaultOnNull<serde_with::OneOrMany<_>>")]
    pub database_id: Vec<i64>,
    #[serde(default)]
    #[serde_as(deserialize_as = "serde_with::DefaultOnNull<serde_with::OneOrMany<_>>")]
    pub toplevel: Vec<bool>,
    #[serde(default)]
    #[serde_as(deserialize_as = "serde_with::DefaultOnNull<serde_with::OneOrMany<_>>")]
    pub calls: Vec<i64>,
}
impl Filter {
    /// Builds new filter with the required fields defined by `group_by` and `timestamp`
    pub fn new(database_id: &[i64]) -> Self {
        Self {
            database_id: database_id.into(),
            ..Self::default()
        }
    }
    fn matches(&self, row: &QueryStat) -> bool {
        (self.database_id.is_empty() || self.database_id.contains(&row.database_id))
            && (self.toplevel.is_empty() || self.toplevel.contains(&row.toplevel))
            && (self.calls.is_empty() || self.calls.contains(&row.calls))
    }
}
#[derive(Clone, Copy, Debug, PartialEq)]
#[doc = concat!(
    " Generated by pco_store to choose which fields to decompress when loading [",
    stringify!(QueryStat), "]"
)]
pub struct Fields {
    database_id: bool,
    toplevel: bool,
    calls: bool,
}
impl Fields {
    pub fn new(fields: &[&str]) -> anyhow::Result<Self> {
        fields.try_into().map_err(|e| anyhow::Error::msg(e))
    }
    pub fn required() -> Self {
        Self {
            database_id: true,
            toplevel: false,
            calls: false,
        }
    }
    fn merge_filter(&mut self, filter: &Filter) {
        (!filter.toplevel.is_empty()).then(|| self.toplevel = true);
        (!filter.calls.is_empty()).then(|| self.calls = true);
    }
    fn select(&self) -> String {
        let mut fields = Vec::new();
        self.database_id.then(|| fields.push("database_id"));
        self.toplevel.then(|| fields.push("toplevel"));
        self.calls.then(|| fields.push("calls"));
        fields.join(", ")
    }
    fn load_from_row(
        &self,
        row: tokio_postgres::Row,
        filter: Option<Filter>,
    ) -> anyhow::Result<CompressedQueryStats> {
        let mut index = 0;
        Ok(CompressedQueryStats {
            filter,
            database_id: if self.database_id {
                let v = row.get(index);
                index += 1;
                v
            } else {
                Default::default()
            },
            toplevel: if self.toplevel {
                let v = row.get(index);
                index += 1;
                v
            } else {
                Default::default()
            },
            calls: if self.calls {
                let v = row.get(index);
                index += 1;
                v
            } else {
                Default::default()
            },
        })
    }
}
impl Default for Fields {
    fn default() -> Self {
        Self {
            database_id: true,
            toplevel: true,
            calls: true,
        }
    }
}
impl TryFrom<&[&str]> for Fields {
    type Error = &'static str;
    fn try_from(input: &[&str]) -> Result<Self, Self::Error> {
        let mut fields = Self::required();
        for s in input {
            match *s {
                "database_id" => fields.database_id = true,
                "toplevel" => fields.toplevel = true,
                "calls" => fields.calls = true,
                _ => return Err("unknown field"),
            }
        }
        Ok(fields)
    }
}
impl<const N: usize> TryFrom<&[&str; N]> for Fields {
    type Error = &'static str;
    fn try_from(input: &[&str; N]) -> Result<Self, Self::Error> {
        Self::try_from(&input[..])
    }
}
impl TryFrom<Vec<String>> for Fields {
    type Error = &'static str;
    fn try_from(input: Vec<String>) -> Result<Self, Self::Error> {
        let input: Vec<_> = input.iter().map(|s| s.as_str()).collect();
        Self::try_from(input.as_slice())
    }
}
impl From<()> for Fields {
    fn from(_: ()) -> Self {
        Self::default()
    }
}
impl<'de> serde::Deserialize<'de> for Fields {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        deserializer.deserialize_any(FieldsVisitor)
    }
}
struct FieldsVisitor;
impl<'de> serde::de::Visitor<'de> for FieldsVisitor {
    type Value = Fields;
    fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
        formatter.write_str("an array of strings matching the struct fields")
    }
    fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
    where
        A: serde::de::SeqAccess<'de>,
    {
        let mut fields = Vec::new();
        while let Some(field) = seq.next_element()? {
            fields.push(field);
        }
        Fields::try_from(fields).map_err(serde::de::Error::custom)
    }
    fn visit_unit<E>(self) -> Result<Self::Value, E>
    where
        E: serde::de::Error,
    {
        Ok(Fields::default())
    }
}
fn serde_compress<T>(items: Vec<T>) -> anyhow::Result<Vec<u8>>
where
    T: serde::Serialize,
{
    use std::io::Write;
    let mut output = Vec::new();
    let mut encoder = zstd::stream::write::Encoder::new(&mut output, 3)?;
    for item in items {
        rmp_serde::encode::write(&mut encoder, &item)?;
    }
    encoder.finish()?;
    Ok(output)
}
fn serde_decompress<T>(input: &[u8]) -> impl Iterator<Item = anyhow::Result<T>> + '_
where
    T: for<'de> serde::Deserialize<'de> + 'static,
{
    let decoder = match zstd::stream::read::Decoder::new(input) {
        Ok(d) => d,
        Err(e) => {
            return Box::new(std::iter::once(Err(e.into())))
                as Box<dyn Iterator<Item = _>>;
        }
    };
    let buffered = std::io::BufReader::with_capacity(128 * 1024, decoder);
    let mut de = rmp_serde::decode::Deserializer::new(buffered);
    Box::new(
        std::iter::from_fn(move || match serde::Deserialize::deserialize(&mut de) {
            Ok(item) => Some(Ok(item)),
            Err(
                rmp_serde::decode::Error::InvalidMarkerRead(ref e),
            ) if e.kind() == std::io::ErrorKind::UnexpectedEof => None,
            Err(e) => Some(Err(e.into())),
        }),
    )
}
fn pco_compress_nested<T>(nested_values: Vec<Vec<T>>) -> anyhow::Result<Vec<u8>>
where
    T: ::pco::data_types::Number,
{
    let mut lengths = Vec::new();
    let mut values = Vec::new();
    for vals in nested_values {
        lengths.push(vals.len() as u64);
        values.extend(vals);
    }
    let length_bytes = ::pco::standalone::simple_compress(
        &lengths,
        &::pco::ChunkConfig::default(),
    )?;
    let value_bytes = ::pco::standalone::simple_compress(
        &values,
        &::pco::ChunkConfig::default(),
    )?;
    let (length_bytes, value_bytes) = (
        serde_bytes::Bytes::new(&length_bytes),
        serde_bytes::Bytes::new(&value_bytes),
    );
    Ok(rmp_serde::to_vec(&(length_bytes, value_bytes))?)
}
fn pco_decompress_nested<T>(bytes: Vec<u8>) -> anyhow::Result<Vec<Vec<T>>>
where
    T: ::pco::data_types::Number,
{
    let (length_bytes, value_bytes): (Vec<u8>, Vec<u8>) = rmp_serde::from_slice(&bytes)?;
    let lengths = ::pco::standalone::simple_decompress::<u64>(&length_bytes)?;
    let values = ::pco::standalone::simple_decompress::<T>(&value_bytes)?;
    let mut values = values.into_iter();
    let mut nested_values = Vec::with_capacity(lengths.len());
    for length in lengths {
        nested_values.push(values.by_ref().take(length as usize).collect::<Vec<T>>());
    }
    Ok(nested_values)
}