Merge pull request #161 from cipherstash/jsonb-contains-variants-fix

feat(ste_vec): add jsonb parameter variants for containment functions

Author: tobyhede

.gitignore

@@ -221,3 +221,11 @@ tests/ste_vec_*M.sql*
 
 # Rust build artifacts (using sccache)
 tests/sqlx/target/
+
+# Work files (agent-generated, not for version control)
+.work/
+.serena/
+
+# Build variants - protect variant deps
+src/deps-protect.txt
+src/deps-ordered-protect.txt

src/ste_vec/functions.sql

@@ -241,18 +241,16 @@ $$ LANGUAGE plpgsql;
 
 
 
---! @brief Extract encrypted JSONB as array for GIN indexing
+--! @brief Extract full encrypted JSONB elements as array
 --!
---! Extracts the encrypted JSONB data and returns it as a native jsonb[]
---! array. This enables efficient GIN indexing using PostgreSQL's built-in array_ops
---! which has native hash support for jsonb elements.
+--! Extracts all JSONB elements from the STE vector including non-deterministic fields.
+--! Use jsonb_array() instead for GIN indexing and containment queries.
 --!
 --! @param val jsonb containing encrypted EQL payload
---! @return jsonb[] Array of JSONB elements for indexing
+--! @return jsonb[] Array of full JSONB elements
 --!
---! @note Preferred for GIN indexes as jsonb has native hash support
---! @see eql_v2.jsonb_array(eql_v2_encrypted)
-CREATE FUNCTION eql_v2.jsonb_array(val jsonb)
+--! @see eql_v2.jsonb_array
+CREATE FUNCTION eql_v2.jsonb_array_from_array_elements(val jsonb)
 RETURNS jsonb[]
 IMMUTABLE STRICT PARALLEL SAFE
 LANGUAGE SQL
@@ -266,21 +264,53 @@ AS $$
 $$;
 
 
---! @brief Extract encrypted JSONB as array from encrypted column value
---!
---! Extracts the encrypted JSONB data from an encrypted column value and returns it as a
---! native jsonb[] array for GIN indexing.
+--! @brief Extract full encrypted JSONB elements as array from encrypted column
 --!
 --! @param val eql_v2_encrypted Encrypted column value
---! @return jsonb[] Array of JSONB elements for indexing
+--! @return jsonb[] Array of full JSONB elements
 --!
---! @example
---! -- Create GIN index for containment queries
---! CREATE INDEX idx_jsonb ON mytable USING GIN (eql_v2.jsonb_array(encrypted_col));
+--! @see eql_v2.jsonb_array_from_array_elements(jsonb)
+CREATE FUNCTION eql_v2.jsonb_array_from_array_elements(val eql_v2_encrypted)
+RETURNS jsonb[]
+IMMUTABLE STRICT PARALLEL SAFE
+LANGUAGE SQL
+AS $$
+  SELECT eql_v2.jsonb_array_from_array_elements(val.data);
+$$;
+
+
+--! @brief Extract deterministic fields as array for GIN indexing
+--!
+--! Extracts only deterministic search term fields (s, b3, hm, ocv, ocf) from each
+--! STE vector element. Excludes non-deterministic ciphertext for correct containment
+--! comparison using PostgreSQL's native @> operator.
+--!
+--! @param val jsonb containing encrypted EQL payload
+--! @return jsonb[] Array of JSONB elements with only deterministic fields
+--!
+--! @note Use this for GIN indexes and containment queries
+--! @see eql_v2.jsonb_contains
+CREATE FUNCTION eql_v2.jsonb_array(val jsonb)
+RETURNS jsonb[]
+IMMUTABLE STRICT PARALLEL SAFE
+LANGUAGE SQL
+AS $$
+  SELECT ARRAY(
+    SELECT jsonb_object_agg(kv.key, kv.value)
+    FROM jsonb_array_elements(
+      CASE WHEN val ? 'sv' THEN val->'sv' ELSE jsonb_build_array(val) END
+    ) AS elem,
+    LATERAL jsonb_each(elem) AS kv(key, value)
+    WHERE kv.key IN ('s', 'b3', 'hm', 'ocv', 'ocf')
+    GROUP BY elem
+  );
+$$;
+
+
+--! @brief Extract deterministic fields as array from encrypted column
 --!
---! -- Query using containment
---! SELECT * FROM mytable
---! WHERE eql_v2.jsonb_array(encrypted_col) @> eql_v2.jsonb_array(search_value);
+--! @param val eql_v2_encrypted Encrypted column value
+--! @return jsonb[] Array of JSONB elements with only deterministic fields
 --!
 --! @see eql_v2.jsonb_array(jsonb)
 CREATE FUNCTION eql_v2.jsonb_array(val eql_v2_encrypted)
@@ -321,6 +351,46 @@ AS $$
 $$;
 
 
+--! @brief GIN-indexable JSONB containment check (encrypted, jsonb)
+--!
+--! Checks if encrypted value 'a' contains all JSONB elements from jsonb value 'b'.
+--! Uses jsonb[] arrays internally for native PostgreSQL GIN index support.
+--!
+--! @param a eql_v2_encrypted Container value (typically a table column)
+--! @param b jsonb JSONB value to search for
+--! @return Boolean True if a contains all elements of b
+--!
+--! @see eql_v2.jsonb_array
+--! @see eql_v2.jsonb_contains(eql_v2_encrypted, eql_v2_encrypted)
+CREATE FUNCTION eql_v2.jsonb_contains(a eql_v2_encrypted, b jsonb)
+RETURNS boolean
+IMMUTABLE STRICT PARALLEL SAFE
+LANGUAGE SQL
+AS $$
+  SELECT eql_v2.jsonb_array(a) @> eql_v2.jsonb_array(b);
+$$;
+
+
+--! @brief GIN-indexable JSONB containment check (jsonb, encrypted)
+--!
+--! Checks if jsonb value 'a' contains all JSONB elements from encrypted value 'b'.
+--! Uses jsonb[] arrays internally for native PostgreSQL GIN index support.
+--!
+--! @param a jsonb Container JSONB value
+--! @param b eql_v2_encrypted Encrypted value to search for
+--! @return Boolean True if a contains all elements of b
+--!
+--! @see eql_v2.jsonb_array
+--! @see eql_v2.jsonb_contains(eql_v2_encrypted, eql_v2_encrypted)
+CREATE FUNCTION eql_v2.jsonb_contains(a jsonb, b eql_v2_encrypted)
+RETURNS boolean
+IMMUTABLE STRICT PARALLEL SAFE
+LANGUAGE SQL
+AS $$
+  SELECT eql_v2.jsonb_array(a) @> eql_v2.jsonb_array(b);
+$$;
+
+
 --! @brief GIN-indexable JSONB "is contained by" check
 --!
 --! Checks if all JSONB elements from 'a' are contained in 'b'.
@@ -341,6 +411,46 @@ AS $$
 $$;
 
 
+--! @brief GIN-indexable JSONB "is contained by" check (encrypted, jsonb)
+--!
+--! Checks if all JSONB elements from encrypted value 'a' are contained in jsonb value 'b'.
+--! Uses jsonb[] arrays internally for native PostgreSQL GIN index support.
+--!
+--! @param a eql_v2_encrypted Value to check (typically a table column)
+--! @param b jsonb Container JSONB value
+--! @return Boolean True if all elements of a are contained in b
+--!
+--! @see eql_v2.jsonb_array
+--! @see eql_v2.jsonb_contained_by(eql_v2_encrypted, eql_v2_encrypted)
+CREATE FUNCTION eql_v2.jsonb_contained_by(a eql_v2_encrypted, b jsonb)
+RETURNS boolean
+IMMUTABLE STRICT PARALLEL SAFE
+LANGUAGE SQL
+AS $$
+  SELECT eql_v2.jsonb_array(a) <@ eql_v2.jsonb_array(b);
+$$;
+
+
+--! @brief GIN-indexable JSONB "is contained by" check (jsonb, encrypted)
+--!
+--! Checks if all JSONB elements from jsonb value 'a' are contained in encrypted value 'b'.
+--! Uses jsonb[] arrays internally for native PostgreSQL GIN index support.
+--!
+--! @param a jsonb Value to check
+--! @param b eql_v2_encrypted Container encrypted value
+--! @return Boolean True if all elements of a are contained in b
+--!
+--! @see eql_v2.jsonb_array
+--! @see eql_v2.jsonb_contained_by(eql_v2_encrypted, eql_v2_encrypted)
+CREATE FUNCTION eql_v2.jsonb_contained_by(a jsonb, b eql_v2_encrypted)
+RETURNS boolean
+IMMUTABLE STRICT PARALLEL SAFE
+LANGUAGE SQL
+AS $$
+  SELECT eql_v2.jsonb_array(a) <@ eql_v2.jsonb_array(b);
+$$;
+
+
 --! @brief Check if STE vector array contains a specific encrypted element
 --!
 --! Tests whether any element in the STE vector array 'a' contains the encrypted value 'b'.

tests/sqlx/src/helpers.rs

@@ -3,6 +3,7 @@
 //! Common utilities for working with encrypted data in tests.
 
 use anyhow::{Context, Result};
+use serde_json;
 use sqlx::{PgPool, Row};
 
 /// Fetch ORE encrypted value from pre-seeded ore table
@@ -83,7 +84,7 @@ pub async fn get_ore_encrypted_as_jsonb(pool: &PgPool, id: i32) -> Result<String
     result.with_context(|| format!("ore table returned NULL for id={}", id))
 }
 
-/// Fetch STE vec encrypted value from a specified table
+/// Fetch STE vec encrypted value from a specified table as serde_json::Value
 ///
 /// Default tables:
 /// - `ste_vec`: Created by migration `003_install_ste_vec_data.sql`, 10 records (ids 1-10)
@@ -93,22 +94,93 @@ pub async fn get_ore_encrypted_as_jsonb(pool: &PgPool, id: i32) -> Result<String
 /// - Records have selectors for $.hello (a7cea93975ed8c01f861ccb6bd082784) with ore_cllw_var_8
 /// - Records have selectors for $.n (2517068c0d1f9d4d41d2c666211f785e) with ore_cllw_u64_8
 ///
+/// Returns the encrypted value as parsed JSON, allowing callers to:
+/// - Inspect structure programmatically
+/// - Use .to_string() when a literal string is needed
+/// - Avoid double-quoting issues with embedded apostrophes
+///
 /// # Arguments
 /// * `pool` - Database connection pool
 /// * `table` - Table name to query (e.g., "ste_vec" or "ste_vec_vast")
 /// * `id` - Row id to fetch
-pub async fn get_ste_vec_encrypted(pool: &PgPool, table: &str, id: i32) -> Result<String> {
-    let sql = format!("SELECT e::text FROM {} WHERE id = {}", table, id);
-    let row = sqlx::query(&sql)
+pub async fn get_ste_vec_encrypted(
+    pool: &PgPool,
+    table: &str,
+    id: i32,
+) -> Result<serde_json::Value> {
+    let sql = format!("SELECT (e).data::jsonb FROM {} WHERE id = {}", table, id);
+    let result: serde_json::Value = sqlx::query_scalar(&sql)
         .fetch_one(pool)
         .await
         .with_context(|| format!("fetching {} encrypted value for id={}", table, id))?;
 
-    let result: Option<String> = row
-        .try_get(0)
-        .with_context(|| format!("extracting text column for id={}", id))?;
+    Ok(result)
+}
+
+/// Fetch two STE vec encrypted values from the same table
+///
+/// Useful for encrypted-to-encrypted containment tests where we need
+/// two distinct encrypted values from the same table.
+///
+/// # Arguments
+/// * `pool` - Database connection pool
+/// * `table` - Table name to query
+/// * `id1` - First row id
+/// * `id2` - Second row id
+///
+/// # Returns
+/// Tuple of (enc1, enc2) as serde_json::Value
+pub async fn get_ste_vec_encrypted_pair(
+    pool: &PgPool,
+    table: &str,
+    id1: i32,
+    id2: i32,
+) -> Result<(serde_json::Value, serde_json::Value)> {
+    let enc1 = get_ste_vec_encrypted(pool, table, id1).await?;
+    let enc2 = get_ste_vec_encrypted(pool, table, id2).await?;
+    Ok((enc1, enc2))
+}
+
+/// Extract a single SV element from an encrypted value as serde_json::Value
+///
+/// Fetches an encrypted value from the specified table and extracts
+/// a specific element from its sv array by index.
+///
+/// # Arguments
+/// * `pool` - Database connection pool
+/// * `table` - Table name to query (e.g., "ste_vec" or "ste_vec_vast")
+/// * `id` - Row id to fetch
+/// * `sv_index` - Index into the sv array (0-based)
+///
+/// # Returns
+/// The sv element as serde_json::Value, suitable for use in containment queries
+/// Use .to_string() when a literal string is needed for SQL interpolation
+pub async fn get_ste_vec_sv_element(
+    pool: &PgPool,
+    table: &str,
+    id: i32,
+    sv_index: i32,
+) -> Result<serde_json::Value> {
+    let sql = format!(
+        "SELECT ((e).data->'sv'->{})::jsonb FROM {} WHERE id = {}",
+        sv_index, table, id
+    );
+    let result: Option<serde_json::Value> = sqlx::query_scalar(&sql)
+        .fetch_one(pool)
+        .await
+        .with_context(|| {
+            format!(
+                "extracting sv element {} from {} id={}",
+                sv_index, table, id
+            )
+        })?;
 
-    result.with_context(|| format!("{} table returned NULL for id={}", table, id))
+    result.with_context(|| {
+        format!(
+            "{} sv element extraction returned NULL for id={}, index={}",
+            table, id, sv_index
+        )
+    })
 }
 
 /// Extract selector term using SQL helper functions

tests/sqlx/src/lib.rs

@@ -13,7 +13,8 @@ pub use assertions::QueryAssertion;
 pub use helpers::{
     analyze_table, assert_uses_index, assert_uses_seq_scan, create_jsonb_gin_index, explain_query,
     get_encrypted_term, get_ore_encrypted, get_ore_encrypted_as_jsonb, get_ste_vec_encrypted,
-    get_ste_vec_selector_term, get_ste_vec_term_by_id,
+    get_ste_vec_encrypted_pair, get_ste_vec_selector_term, get_ste_vec_sv_element,
+    get_ste_vec_term_by_id,
 };
 pub use index_types as IndexTypes;
 pub use selectors::Selectors;