fix(sqlite): deserialize JSON columns (0xPlaygrounds#1797)

Author: gold-silver-copper

crates/rig-sqlite/src/lib.rs

@@ -1928,12 +1928,50 @@ fn sqlite_text_value(
     Ok(serde_json::Value::String(value.to_string()))
 }
 
+fn sqlite_column_declares_json(column_type: &str) -> bool {
+    column_type
+        .split_whitespace()
+        .next()
+        .is_some_and(|token| token.eq_ignore_ascii_case("JSON"))
+}
+
+fn sqlite_json_text_value(
+    index: usize,
+    value_type: Type,
+    column: &Column,
+    value: &[u8],
+) -> rusqlite::Result<serde_json::Value> {
+    let value = std::str::from_utf8(value).map_err(|e| {
+        sqlite_column_value_error(
+            index,
+            value_type,
+            column,
+            format!("invalid UTF-8 JSON text: {e}"),
+        )
+    })?;
+
+    serde_json::from_str(value).map_err(|e| {
+        sqlite_column_value_error(index, value_type, column, format!("invalid JSON text: {e}"))
+    })
+}
+
 fn sqlite_column_value_to_json(
     index: usize,
     column: &Column,
     value: ValueRef<'_>,
 ) -> rusqlite::Result<serde_json::Value> {
     let value_type = value.data_type();
+
+    if sqlite_column_declares_json(column.col_type) {
+        return match value {
+            ValueRef::Null => Ok(serde_json::Value::Null),
+            ValueRef::Text(value) => sqlite_json_text_value(index, value_type, column, value),
+            ValueRef::Integer(value) => Ok(serde_json::Value::Number(value.into())),
+            ValueRef::Real(value) => sqlite_number_value(index, value_type, column, value),
+            ValueRef::Blob(value) => sqlite_json_text_value(index, value_type, column, value),
+        };
+    }
+
     let column_affinity = SqliteColumnAffinity::from_column_type(column.col_type);
 
     match (column_affinity, value) {
@@ -2273,6 +2311,66 @@ mod tests {
         ]
     }
 
+    #[test]
+    fn json_column_text_decodes_to_json_object() -> anyhow::Result<()> {
+        let column = Column::new("metadata", "JSON");
+        let value = sqlite_column_value_to_json(
+            0,
+            &column,
+            ValueRef::Text(br#"{"knowledge_doc_id":361,"knowledge_id":1,"user_id":1}"#),
+        )?;
+
+        let expected = serde_json::json!({
+            "knowledge_doc_id": 361,
+            "knowledge_id": 1,
+            "user_id": 1
+        });
+        anyhow::ensure!(
+            value == expected,
+            "JSON column text should decode to a JSON object, got {value:?}"
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn text_column_json_looking_text_stays_string() -> anyhow::Result<()> {
+        let column = Column::new("metadata", "TEXT");
+        let value = sqlite_column_value_to_json(
+            0,
+            &column,
+            ValueRef::Text(br#"{"knowledge_doc_id":361,"knowledge_id":1,"user_id":1}"#),
+        )?;
+
+        let expected =
+            serde_json::json!(r#"{"knowledge_doc_id":361,"knowledge_id":1,"user_id":1}"#);
+        anyhow::ensure!(
+            value == expected,
+            "TEXT column should preserve JSON-looking text as a string, got {value:?}"
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn json_column_invalid_text_returns_conversion_error() -> anyhow::Result<()> {
+        let column = Column::new("metadata", "JSON");
+        let err = match sqlite_column_value_to_json(0, &column, ValueRef::Text(b"not json")) {
+            Ok(value) => anyhow::bail!("invalid JSON column text should fail, got {value:?}"),
+            Err(err) => err,
+        };
+
+        anyhow::ensure!(
+            matches!(
+                err,
+                rusqlite::Error::FromSqlConversionFailure(0, Type::Text, _)
+            ),
+            "invalid JSON column text should return a conversion error, got {err}"
+        );
+
+        Ok(())
+    }
+
     fn filter_error<T: std::fmt::Debug>(
         result: Result<T, FilterError>,
         context: &str,
@@ -3230,6 +3328,50 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn live_json_column_structured_metadata_round_trips_in_top_n() -> anyhow::Result<()> {
+        let metadata = StructuredMetadata {
+            user_id: 1,
+            knowledge_id: 1,
+            knowledge_doc_id: 361,
+        };
+        let index = live_structured_json_metadata_test_index(
+            "live_json_column_structured_metadata_round_trips_in_top_n",
+            vec![structured_json_metadata_row(
+                "structured",
+                metadata.clone(),
+                "metadata document",
+                vec![1.0, 0.0],
+            )],
+        )
+        .await?;
+
+        let req = VectorSearchRequest::<SqliteSearchFilter>::builder()
+            .query("needle")
+            .samples(1)
+            .build();
+        let results = index
+            .top_n::<StructuredJsonMetadataDocument>(req.clone())
+            .await?;
+
+        let Some((_, id, doc)) = results.first() else {
+            anyhow::bail!("expected structured JSON metadata document result");
+        };
+        anyhow::ensure!(id == "structured", "unexpected id: {id}");
+        anyhow::ensure!(
+            doc.metadata == metadata,
+            "JSON column should deserialize into structured metadata: {doc:?}"
+        );
+
+        let id_results = index.top_n_ids(req).await?;
+        anyhow::ensure!(
+            id_results.first().is_some_and(|(_, id)| id == "structured"),
+            "top_n_ids should still return the structured metadata document id: {id_results:?}"
+        );
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn live_text_affinity_metadata_filters_during_candidate_search() -> anyhow::Result<()> {
         let index = live_common_type_test_index(
@@ -4140,6 +4282,22 @@ mod tests {
         Ok(vector_store.index(model))
     }
 
+    async fn live_structured_json_metadata_test_index(
+        name: &str,
+        rows: Vec<(StructuredJsonMetadataDocument, OneOrMany<Embedding>)>,
+    ) -> anyhow::Result<SqliteVectorIndex<TestEmbeddingModel, StructuredJsonMetadataDocument>> {
+        register_sqlite_vec_extension();
+
+        let conn = Connection::open(format!("file:{name}?mode=memory")).await?;
+        let model = TestEmbeddingModel;
+        let vector_store: SqliteVectorStore<_, StructuredJsonMetadataDocument> =
+            SqliteVectorStore::new(conn, &model).await?;
+
+        vector_store.add_rows(rows).await?;
+
+        Ok(vector_store.index(model))
+    }
+
     fn row(
         id: impl Into<String>,
         category: impl Into<String>,
@@ -4207,6 +4365,27 @@ mod tests {
         )
     }
 
+    fn structured_json_metadata_row(
+        id: impl Into<String>,
+        metadata: StructuredMetadata,
+        title: impl Into<String>,
+        vec: Vec<f64>,
+    ) -> (StructuredJsonMetadataDocument, OneOrMany<Embedding>) {
+        let document = StructuredJsonMetadataDocument {
+            id: id.into(),
+            metadata,
+            title: title.into(),
+        };
+
+        (
+            document.clone(),
+            OneOrMany::one(Embedding {
+                document: document.title.clone(),
+                vec,
+            }),
+        )
+    }
+
     fn reordered_id_row(
         id: impl Into<String>,
         title: impl Into<String>,
@@ -4675,6 +4854,49 @@ mod tests {
         }
     }
 
+    #[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
+    struct StructuredMetadata {
+        user_id: i64,
+        knowledge_id: i64,
+        knowledge_doc_id: i64,
+    }
+
+    #[derive(Clone, Debug, Deserialize, Serialize)]
+    struct StructuredJsonMetadataDocument {
+        id: String,
+        metadata: StructuredMetadata,
+        title: String,
+    }
+
+    impl SqliteVectorStoreTable for StructuredJsonMetadataDocument {
+        fn name() -> &'static str {
+            "live_structured_json_metadata_test_documents"
+        }
+
+        fn schema() -> Vec<Column> {
+            vec![
+                Column::new("id", "TEXT PRIMARY KEY"),
+                Column::new("metadata", "JSON"),
+                Column::new("title", "TEXT"),
+            ]
+        }
+
+        fn id(&self) -> String {
+            self.id.clone()
+        }
+
+        fn column_values(&self) -> Vec<(&'static str, Box<dyn ColumnValue>)> {
+            vec![
+                ("id", Box::new(self.id.clone())),
+                (
+                    "metadata",
+                    Box::new(serde_json::to_string(&self.metadata).unwrap_or_default()),
+                ),
+                ("title", Box::new(self.title.clone())),
+            ]
+        }
+    }
+
     #[derive(Clone, Debug, Deserialize, Serialize)]
     struct TypedTestDocument {
         id: i64,