Support filtering by any fields

Author: seanlinsley

Cargo.lock

@@ -20,6 +20,9 @@ futures = "0.3"
 pco = "0.4"
 proc-macro2 = "1.0"
 quote = "1.0"
+serde = { version = "1.0", features = ["derive"] }
+serde_json = { version = "1.0" }
+serde_with = "3.16"
 syn = "2.0"
 tokio-postgres = { version = "0.7", features = ["with-chrono-0_4"] }
 

Cargo.toml

@@ -71,76 +71,60 @@ pub fn store(args: TokenStream, item: TokenStream) -> TokenStream {
     };
 
     // load and delete
-    let mut fields = Vec::new();
-    let mut load_filters = Vec::new();
+    let mut fields = vec![quote!{ pub filter: Option<&'a Filter>, }];
     let mut load_checks = Vec::new();
     let mut load_where = Vec::new();
     let mut load_params = Vec::new();
     let mut load_fields = Vec::new();
     let mut bind = 1;
     let mut index: usize = 0;
-    let mut timestamp_ty: Option<_> = None;
+    let mut timestamp_ty = None;
     let mut using_chrono = false;
     for field in model.fields.iter() {
         let ident = field.ident.clone().unwrap();
         let ty = field.ty.clone();
+        let name = format!("{ident}");
         if group_by.iter().any(|i| *i == ident) {
             fields.push(quote! { pub #ident: #ty, });
-            load_filters.push(quote! { #ident: &[#ty], });
-            let name = format!("{ident}");
             load_checks.push(quote! {
-                if #ident.is_empty() {
-                    return Err(anyhow::Error::msg(#name.to_string() + "must be specified"));
+                if filter.#ident.is_empty() {
+                    return Err(anyhow::Error::msg(#name.to_string() + " is required"));
                 }
             });
             load_where.push(format!("{ident} = ANY(${bind})"));
             bind += 1;
-            load_params.push(quote! { &#ident, });
+            load_params.push(quote! { &filter.#ident, });
         } else if timestamp.as_ref().map(|t| *t == ident).unwrap_or(false) {
-            timestamp_ty = Some(ty.clone());
+            timestamp_ty = Some(quote! { #ty });
             using_chrono = !ty.to_token_stream().to_string().contains("SystemTime");
-            fields.push(quote! {
-                pub filter: bool,
-                pub filter_start: #ty,
-                pub filter_end: #ty,
-                pub start_at: #ty,
-                pub end_at: #ty,
-                #ident: Vec<u8>,
+            fields.push(quote! { #ident: Vec<u8>, });
+            load_checks.push(quote! {
+                if filter.#ident.is_none() {
+                    return Err(anyhow::Error::msg(#name.to_string() + " is required"));
+                }
             });
             load_where.push(format!("end_at >= ${bind}"));
             bind += 1;
             load_where.push(format!("start_at <= ${bind}"));
             bind += 1;
-            load_params.push(quote! { &filter_start, &filter_end, });
+            load_params.push(quote! {
+                filter.#ident.as_ref().unwrap().start(),
+                filter.#ident.as_ref().unwrap().end(),
+            });
         } else {
             fields.push(quote! { #ident: Vec<u8>, });
         }
         if timestamp.as_ref().map(|t| *t == ident).unwrap_or(false) {
-            load_fields.push(quote! { start_at: row.get(#index), });
-            index += 1;
-            load_fields.push(quote! { end_at: row.get(#index), });
-            index += 1;
+            index += 2; // Skip over start_at and end_at
         }
         load_fields.push(quote! { #ident: row.get(#index), });
         index += 1;
     }
-    let mut delete_fields = load_fields.clone();
-    if timestamp.is_some() {
-        let ty = timestamp_ty.unwrap_or_else(|| Type::Verbatim(quote! { std::time::SystemTime }));
-        load_filters.push(quote! {
-            filter_start: #ty,
-            filter_end: #ty,
-        });
-        load_fields.push(quote! { filter: true, filter_start, filter_end, });
-        delete_fields.push(quote! { filter: false, filter_start, filter_end, });
-    }
     let fields = tokens(fields);
-    let load_filters = tokens(load_filters);
     let load_checks = tokens(load_checks);
     let load_where = if load_where.is_empty() { "true".to_string() } else { load_where.join(" AND ") };
     let load_params = tokens(load_params);
     let load_fields = tokens(load_fields);
-    let delete_fields = tokens(delete_fields);
     let load_sql = format!("SELECT * FROM {table_name} WHERE {load_where}");
     let delete_sql = format!("DELETE FROM {table_name} WHERE {load_where} RETURNING *");
 
@@ -204,11 +188,6 @@ pub fn store(args: TokenStream, item: TokenStream) -> TokenStream {
     let decompress_fields = tokens(decompress_fields);
     let compressed_field_sizes = tokens(compressed_field_sizes);
     let decompressed_fields = tokens(decompressed_fields);
-    let in_time_range = if timestamp.is_some() {
-        quote! { !self.filter || row.#timestamp >= self.filter_start && row.#timestamp <= self.filter_end }
-    } else {
-        quote! { true }
-    };
 
     // store
     let mut store_fields = Vec::new();
@@ -281,39 +260,67 @@ pub fn store(args: TokenStream, item: TokenStream) -> TokenStream {
     };
     let store_sql = format!("COPY {table_name} ({store_fields}) FROM STDIN BINARY");
 
+    let mut filter_fields = Vec::new();
+    let mut filter_conditions = Vec::new();
+    for field in model.fields.iter() {
+        let ident = field.ident.clone().unwrap();
+        let ty = field.ty.clone();
+        let time = ty.to_token_stream().to_string().contains("Time");
+        if time {
+            filter_fields.push(quote! {
+                #[serde(deserialize_with = "serde_extra::deserialize_time_range")]
+                pub #ident: Option<std::ops::RangeInclusive<#ty>>,
+            });
+        } else {
+            filter_fields.push(quote! {
+                #[serde(default)]
+                #[serde_as(deserialize_as = "serde_with::OneOrMany<_>")]
+                pub #ident: Vec<#ty>,
+            });
+        }
+        if time {
+            filter_conditions.push(quote! {
+                self.#ident.as_ref().map(|t| t.contains(&row.#ident)) != Some(false)
+            });
+        } else {
+            filter_conditions.push(quote! {
+                (self.#ident.is_empty() || self.#ident.contains(&row.#ident))
+            });
+        }
+    }
+    let filter_fields = tokens(filter_fields);
+
+    let serde_extra = timestamp_ty.map(|t| serde_extra(&t));
+
     quote! {
         #item
 
         #[doc=concat!(" Generated by pco_store to store and load compressed versions of [", stringify!(#name), "]")]
-        pub struct #packed_name {
+        pub struct #packed_name<'a> {
             #fields
         }
 
-        impl #packed_name {
+        impl<'a> #packed_name<'a> {
             /// Loads data for the specified filters.
-            ///
-            /// For models with a timestamp, [decompress][Self::decompress] automatically filters out
-            /// rows outside the requested time range.
-            pub async fn load(db: &deadpool_postgres::Object, #load_filters) -> anyhow::Result<Vec<#packed_name>> {
+            pub async fn load(db: &deadpool_postgres::Object, filter: &'a Filter) -> anyhow::Result<Vec<#packed_name<'a>>> {
                 #load_checks
                 let sql = #load_sql;
                 let mut results = Vec::new();
                 for row in db.query(&db.prepare_cached(&sql).await?, &[#load_params]).await? {
-                    results.push(#packed_name { #load_fields });
+                    results.push(#packed_name { filter: Some(filter), #load_fields });
                 }
                 Ok(results)
             }
 
             /// Deletes data for the specified filters, returning it to the caller.
             ///
-            /// For models with a timestamp, [decompress][Self::decompress] **will not** filter out
-            /// rows outside the requested time range.
-            pub async fn delete(db: &deadpool_postgres::Object, #load_filters) -> anyhow::Result<Vec<#packed_name>> {
+            /// Note that all rows are returned from [decompress][Self::decompress] even if post-decompress filters would normally apply.
+            pub async fn delete(db: &deadpool_postgres::Object, filter: &'a Filter) -> anyhow::Result<Vec<#packed_name<'a>>> {
                 #load_checks
                 let sql = #delete_sql;
                 let mut results = Vec::new();
                 for row in db.query(&db.prepare_cached(&sql).await?, &[#load_params]).await? {
-                    results.push(#packed_name { #delete_fields });
+                    results.push(#packed_name { filter: None, #load_fields });
                 }
                 Ok(results)
             }
@@ -325,7 +332,7 @@ pub fn store(args: TokenStream, item: TokenStream) -> TokenStream {
                 let len = [#compressed_field_sizes].into_iter().max().unwrap_or(0);
                 for index in 0..len {
                     let row = #name { #decompressed_fields };
-                    if #in_time_range {
+                    if self.filter.as_ref().map(|f| f.filter(&row)) != Some(false) {
                         results.push(row);
                     }
                 }
@@ -387,10 +394,90 @@ pub fn store(args: TokenStream, item: TokenStream) -> TokenStream {
                 Ok(())
             }
         }
+
+        #[serde_with::serde_as]
+        #[derive(Debug, Default, serde::Deserialize, Clone)]
+        #[serde(deny_unknown_fields)]
+        #[doc=concat!(" Generated by pco_store to specify filters when loading [", stringify!(#name), "]")]
+        pub struct Filter {
+            #filter_fields
+        }
+
+        impl Filter {
+            pub fn filter(&self, row: &#name) -> bool {
+                #(#filter_conditions)&&*
+            }
+        }
+
+        #serde_extra
     }
     .into()
 }
 
+fn serde_extra(timestamp_ty: &proc_macro2::TokenStream) -> proc_macro2::TokenStream {
+    quote! {
+        mod serde_extra {
+            use super::*;
+            use serde::de::{self, SeqAccess, Visitor};
+            use serde::{Deserialize, Deserializer};
+            use std::fmt;
+            use std::ops::RangeInclusive;
+
+            pub(super) fn deserialize_time_range<'de, D>(deserializer: D) -> Result<Option<RangeInclusive<#timestamp_ty>>, D::Error>
+            where
+                D: Deserializer<'de>,
+            {
+                Ok(TimeRange::deserialize(deserializer)?.0)
+            }
+
+            #[derive(Debug, PartialEq)]
+            struct TimeRange(Option<RangeInclusive<#timestamp_ty>>);
+            impl<'de> Deserialize<'de> for TimeRange {
+                fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+                where
+                    D: Deserializer<'de>,
+                {
+                    deserializer.deserialize_any(TimeRangeVisitor)
+                }
+            }
+
+            struct TimeRangeVisitor;
+            impl<'de> Visitor<'de> for TimeRangeVisitor {
+                type Value = TimeRange;
+
+                fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+                    formatter.write_str("a single time string, or an array with 1-2 time strings")
+                }
+
+                fn visit_str<E>(self, value: &str) -> Result<Self::Value, E>
+                where
+                    E: de::Error,
+                {
+                    match serde_json::from_str::<#timestamp_ty>(value) {
+                        Ok(start) => Ok(TimeRange(Some(start..=start))),
+                        Err(err) => Err(E::custom(format!("invalid time format: {err}"))),
+                    }
+                }
+
+                fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
+                where
+                    A: SeqAccess<'de>,
+                {
+                    let start = match seq.next_element::<Option<#timestamp_ty>>()? {
+                        Some(Some(time)) => time,
+                        Some(None) | None => return Ok(TimeRange(None)),
+                    };
+                    let end = match seq.next_element::<Option<#timestamp_ty>>()? {
+                        Some(Some(time)) => time,
+                        Some(None) | None => start,
+                    };
+                    Ok(TimeRange(Some(start..=end)))
+                }
+            }
+        }
+    }
+}
+
 fn copy_type(rust_type: String) -> &'static str {
     match rust_type.as_str() {
         "f32" => "FLOAT4",