From 064b973a2ac346e28358a82bc5e8251f5d106401 Mon Sep 17 00:00:00 2001
From: Abdo <abdo@abdnh.net>
Date: Fri, 1 Sep 2023 07:13:31 +0300
Subject: [PATCH] Allow searching for missing custom data properties (#2626)

* Allow searching for missing custom data properties

* Add has-cd query
---
 rslib/src/search/parser.rs    |  2 ++
 rslib/src/search/sqlwriter.rs |  8 ++++++++
 rslib/src/search/writer.rs    |  1 +
 rslib/src/storage/sqlite.rs   | 23 +++++++++++++++++++++++
 4 files changed, 34 insertions(+)

diff --git a/rslib/src/search/parser.rs b/rslib/src/search/parser.rs
index 8c80da2f3..ddb5b1249 100644
--- a/rslib/src/search/parser.rs
+++ b/rslib/src/search/parser.rs
@@ -92,6 +92,7 @@ pub enum SearchNode {
     Regex(String),
     NoCombining(String),
     WordBoundary(String),
+    CustomData(String),
 }
 
 #[derive(Debug, PartialEq, Clone)]
@@ -348,6 +349,7 @@ fn search_node_for_text_with_argument<'a>(
         "nc" => SearchNode::NoCombining(unescape(val)?),
         "w" => SearchNode::WordBoundary(unescape(val)?),
         "dupe" => parse_dupe(val)?,
+        "has-cd" => SearchNode::CustomData(unescape(val)?),
         // anything else is a field search
         _ => parse_single_field(key, val)?,
     })
diff --git a/rslib/src/search/sqlwriter.rs b/rslib/src/search/sqlwriter.rs
index 89463e991..42715a0c5 100644
--- a/rslib/src/search/sqlwriter.rs
+++ b/rslib/src/search/sqlwriter.rs
@@ -179,6 +179,7 @@ impl SqlWriter<'_> {
                 write!(self.sql, "c.id in ({})", cids).unwrap();
             }
             SearchNode::Property { operator, kind } => self.write_prop(operator, kind)?,
+            SearchNode::CustomData(key) => self.write_custom_data(key)?,
             SearchNode::WholeCollection => write!(self.sql, "true").unwrap(),
         };
         Ok(())
@@ -368,6 +369,12 @@ impl SqlWriter<'_> {
         Ok(())
     }
 
+    fn write_custom_data(&mut self, key: &str) -> Result<()> {
+        write!(self.sql, "has_custom_data(c.data, '{key}')").unwrap();
+
+        Ok(())
+    }
+
     fn write_state(&mut self, state: &StateKind) -> Result<()> {
         let timing = self.col.timing_today()?;
         match state {
@@ -927,6 +934,7 @@ impl SearchNode {
             SearchNode::Flag(_) => RequiredTable::Cards,
             SearchNode::CardIds(_) => RequiredTable::Cards,
             SearchNode::Property { .. } => RequiredTable::Cards,
+            SearchNode::CustomData { .. } => RequiredTable::Cards,
 
             SearchNode::UnqualifiedText(_) => RequiredTable::Notes,
             SearchNode::SingleField { .. } => RequiredTable::Notes,
diff --git a/rslib/src/search/writer.rs b/rslib/src/search/writer.rs
index b43522cfa..b1445cde2 100644
--- a/rslib/src/search/writer.rs
+++ b/rslib/src/search/writer.rs
@@ -86,6 +86,7 @@ fn write_search_node(node: &SearchNode) -> String {
         Regex(s) => maybe_quote(&format!("re:{}", s)),
         NoCombining(s) => maybe_quote(&format!("nc:{}", s)),
         WordBoundary(s) => maybe_quote(&format!("w:{}", s)),
+        CustomData(k) => maybe_quote(&format!("has-cd:{}", k)),
     }
 }
 
diff --git a/rslib/src/storage/sqlite.rs b/rslib/src/storage/sqlite.rs
index ce48ce779..ff235c825 100644
--- a/rslib/src/storage/sqlite.rs
+++ b/rslib/src/storage/sqlite.rs
@@ -70,6 +70,7 @@ fn open_or_create_collection_db(path: &Path) -> Result<Connection> {
     add_without_combining_function(&db)?;
     add_fnvhash_function(&db)?;
     add_extract_custom_data_number_function(&db)?;
+    add_has_custom_data_function(&db)?;
 
     db.create_collation("unicase", unicase_compare)?;
 
@@ -223,6 +224,28 @@ fn add_extract_custom_data_number_function(db: &Connection) -> rusqlite::Result<
     )
 }
 
+/// eg. has_custom_data(card.data, 'r') -> bool
+fn add_has_custom_data_function(db: &Connection) -> rusqlite::Result<()> {
+    db.create_scalar_function(
+        "has_custom_data",
+        2,
+        FunctionFlags::SQLITE_DETERMINISTIC,
+        move |ctx| {
+            assert_eq!(ctx.len(), 2, "called with unexpected number of arguments");
+
+            let Ok(card_data) = ctx.get_raw(0).as_str() else { return Ok(None) };
+            if card_data.is_empty() {
+                return Ok(Some(false));
+            }
+            let Ok(key) = ctx.get_raw(1).as_str() else { return Ok(Some(false)) };
+            let custom_data = &CardData::from_str(card_data).custom_data;
+            let Ok(value) = serde_json::from_str::<Value>(custom_data) else { return Ok(Some(false)) };
+
+            Ok(value.get(key).map(|_| true))
+        },
+    )
+}
+
 /// Fetch schema version from database.
 /// Return (must_create, version)
 fn schema_version(db: &Connection) -> Result<(bool, u8)> {