prax_pgvector/
query.rs

1//! High-level query builder for vector similarity search.
2//!
3//! This module provides a fluent builder API for constructing vector search queries
4//! that integrate with the prax-postgres engine.
5//!
6//! # Examples
7//!
8//! ```rust
9//! use prax_pgvector::query::VectorSearchBuilder;
10//! use prax_pgvector::{Embedding, DistanceMetric};
11//!
12//! let query = VectorSearchBuilder::new("documents", "embedding")
13//!     .query(Embedding::new(vec![0.1, 0.2, 0.3]))
14//!     .metric(DistanceMetric::Cosine)
15//!     .limit(10)
16//!     .select(&["id", "title", "content"])
17//!     .where_clause("category = 'tech'")
18//!     .build();
19//!
20//! let sql = query.to_sql();
21//! assert!(sql.contains("<=>")); // cosine distance operator
22//! ```
23
24use serde::{Deserialize, Serialize};
25
26use prax_query::sql::is_valid_sql_identifier;
27
28use crate::error::{VectorError, VectorResult};
29use crate::ops::{DistanceMetric, SearchParams};
30use crate::types::Embedding;
31
32/// A fully constructed vector search query ready for execution.
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct VectorSearchQuery {
35    /// The table to search.
36    pub table: String,
37    /// The vector column.
38    pub column: String,
39    /// The query vector.
40    pub query_vector: Embedding,
41    /// Distance metric.
42    pub metric: DistanceMetric,
43    /// Maximum number of results.
44    pub limit: usize,
45    /// Columns to select (empty = all).
46    pub select_columns: Vec<String>,
47    /// Additional WHERE conditions.
48    pub where_clauses: Vec<String>,
49    /// Whether to include the distance in results.
50    pub include_distance: bool,
51    /// Alias for the distance column.
52    pub distance_alias: String,
53    /// Maximum distance threshold (radius search).
54    pub max_distance: Option<f64>,
55    /// Minimum distance threshold.
56    pub min_distance: Option<f64>,
57    /// Additional ORDER BY clauses (after distance).
58    pub extra_order_by: Vec<String>,
59    /// Offset for pagination.
60    pub offset: Option<usize>,
61    /// Search parameters (probes, ef_search).
62    pub search_params: SearchParams,
63}
64
65impl VectorSearchQuery {
66    /// Generate the complete SQL query.
67    ///
68    /// The query vector should be passed as parameter `$1`.
69    pub fn to_sql(&self) -> String {
70        self.to_sql_with_param(1)
71    }
72
73    /// Generate the complete SQL query with a custom parameter index.
74    pub fn to_sql_with_param(&self, param_index: usize) -> String {
75        let param = format!("${param_index}");
76        let distance_expr = format!("{} {} {}", self.column, self.metric.operator(), param);
77
78        // SELECT clause
79        let select = if self.select_columns.is_empty() {
80            "*".to_string()
81        } else {
82            self.select_columns.join(", ")
83        };
84
85        let distance_select = if self.include_distance {
86            format!(", {} AS {}", distance_expr, self.distance_alias)
87        } else {
88            String::new()
89        };
90
91        // WHERE clause
92        let mut where_parts = Vec::new();
93
94        if let Some(max) = self.max_distance {
95            where_parts.push(format!("{distance_expr} < {max}"));
96        }
97        if let Some(min) = self.min_distance {
98            where_parts.push(format!("{distance_expr} >= {min}"));
99        }
100        where_parts.extend(self.where_clauses.clone());
101
102        let where_clause = if where_parts.is_empty() {
103            String::new()
104        } else {
105            format!(" WHERE {}", where_parts.join(" AND "))
106        };
107
108        // ORDER BY clause
109        let order_by_main = if self.include_distance {
110            self.distance_alias.clone()
111        } else {
112            distance_expr
113        };
114
115        let order_by = if self.extra_order_by.is_empty() {
116            order_by_main
117        } else {
118            let mut parts = vec![order_by_main];
119            parts.extend(self.extra_order_by.clone());
120            parts.join(", ")
121        };
122
123        // LIMIT and OFFSET
124        let limit = format!(" LIMIT {}", self.limit);
125        let offset = self
126            .offset
127            .map(|o| format!(" OFFSET {o}"))
128            .unwrap_or_default();
129
130        format!(
131            "SELECT {}{} FROM {}{}  ORDER BY {}{}{}",
132            select, distance_select, self.table, where_clause, order_by, limit, offset
133        )
134    }
135
136    /// Generate SET commands for search parameters.
137    ///
138    /// These should be executed before the search query to tune index scan behavior.
139    pub fn param_set_sql(&self) -> Vec<String> {
140        self.search_params.to_set_sql()
141    }
142}
143
144/// Fluent builder for vector search queries.
145///
146/// # Examples
147///
148/// ```rust
149/// use prax_pgvector::query::VectorSearchBuilder;
150/// use prax_pgvector::{Embedding, DistanceMetric};
151///
152/// let query = VectorSearchBuilder::new("documents", "embedding")
153///     .query(Embedding::new(vec![0.1, 0.2, 0.3]))
154///     .metric(DistanceMetric::Cosine)
155///     .limit(10)
156///     .ef_search(200)
157///     .build();
158/// ```
159pub struct VectorSearchBuilder {
160    table: String,
161    column: String,
162    query_vector: Option<Embedding>,
163    metric: DistanceMetric,
164    limit: usize,
165    select_columns: Vec<String>,
166    where_clauses: Vec<String>,
167    include_distance: bool,
168    distance_alias: String,
169    max_distance: Option<f64>,
170    min_distance: Option<f64>,
171    extra_order_by: Vec<String>,
172    offset: Option<usize>,
173    search_params: SearchParams,
174}
175
176impl VectorSearchBuilder {
177    /// Create a new search builder for a table and vector column.
178    pub fn new(table: impl Into<String>, column: impl Into<String>) -> Self {
179        Self {
180            table: table.into(),
181            column: column.into(),
182            query_vector: None,
183            metric: DistanceMetric::L2,
184            limit: 10,
185            select_columns: Vec::new(),
186            where_clauses: Vec::new(),
187            include_distance: true,
188            distance_alias: "distance".to_string(),
189            max_distance: None,
190            min_distance: None,
191            extra_order_by: Vec::new(),
192            offset: None,
193            search_params: SearchParams::new(),
194        }
195    }
196
197    /// Set the query vector.
198    pub fn query(mut self, embedding: Embedding) -> Self {
199        self.query_vector = Some(embedding);
200        self
201    }
202
203    /// Set the distance metric.
204    pub fn metric(mut self, metric: DistanceMetric) -> Self {
205        self.metric = metric;
206        self
207    }
208
209    /// Set the result limit.
210    pub fn limit(mut self, limit: usize) -> Self {
211        self.limit = limit;
212        self
213    }
214
215    /// Set specific columns to select.
216    pub fn select(mut self, columns: &[&str]) -> Self {
217        self.select_columns = columns.iter().map(|c| (*c).to_string()).collect();
218        self
219    }
220
221    /// Add a WHERE condition.
222    pub fn where_clause(mut self, condition: impl Into<String>) -> Self {
223        self.where_clauses.push(condition.into());
224        self
225    }
226
227    /// Set the maximum distance (radius search).
228    pub fn max_distance(mut self, distance: f64) -> Self {
229        self.max_distance = Some(distance);
230        self
231    }
232
233    /// Set the minimum distance.
234    pub fn min_distance(mut self, distance: f64) -> Self {
235        self.min_distance = Some(distance);
236        self
237    }
238
239    /// Don't include the distance in the results.
240    pub fn without_distance(mut self) -> Self {
241        self.include_distance = false;
242        self
243    }
244
245    /// Set a custom distance column alias.
246    pub fn distance_alias(mut self, alias: impl Into<String>) -> Self {
247        self.distance_alias = alias.into();
248        self
249    }
250
251    /// Add an additional ORDER BY clause (after distance).
252    pub fn then_order_by(mut self, clause: impl Into<String>) -> Self {
253        self.extra_order_by.push(clause.into());
254        self
255    }
256
257    /// Set the offset for pagination.
258    pub fn offset(mut self, offset: usize) -> Self {
259        self.offset = Some(offset);
260        self
261    }
262
263    /// Set the IVFFlat probes parameter.
264    pub fn probes(mut self, probes: usize) -> Self {
265        self.search_params = self.search_params.probes(probes);
266        self
267    }
268
269    /// Set the HNSW ef_search parameter.
270    pub fn ef_search(mut self, ef: usize) -> Self {
271        self.search_params = self.search_params.ef_search(ef);
272        self
273    }
274
275    /// Build the vector search query.
276    ///
277    /// # Panics
278    ///
279    /// Panics if no query vector has been set. Use [`Self::try_build`] for
280    /// a non-panicking alternative.
281    pub fn build(self) -> VectorSearchQuery {
282        self.try_build()
283            .expect("query vector must be set before building")
284    }
285
286    /// Try to build the vector search query.
287    ///
288    /// Returns `None` if no query vector has been set.
289    pub fn try_build(self) -> Option<VectorSearchQuery> {
290        let query_vector = self.query_vector?;
291
292        Some(VectorSearchQuery {
293            table: self.table,
294            column: self.column,
295            query_vector,
296            metric: self.metric,
297            limit: self.limit,
298            select_columns: self.select_columns,
299            where_clauses: self.where_clauses,
300            include_distance: self.include_distance,
301            distance_alias: self.distance_alias,
302            max_distance: self.max_distance,
303            min_distance: self.min_distance,
304            extra_order_by: self.extra_order_by,
305            offset: self.offset,
306            search_params: self.search_params,
307        })
308    }
309}
310
311/// Returns `true` if `name` is a valid PostgreSQL text search configuration
312/// name — a simple identifier matching `^[A-Za-z_][A-Za-z0-9_]*$`.
313fn is_valid_ts_config_name(name: &str) -> bool {
314    is_valid_sql_identifier(name)
315}
316
317/// Builder for hybrid search queries that combine vector similarity with full-text search.
318///
319/// This generates queries that use both pgvector distance operators and
320/// PostgreSQL tsvector/tsquery for combined similarity scoring.
321///
322/// # Examples
323///
324/// ```rust
325/// use prax_pgvector::query::HybridSearchBuilder;
326/// use prax_pgvector::{Embedding, DistanceMetric};
327///
328/// let query = HybridSearchBuilder::new("documents")
329///     .vector_column("embedding")
330///     .text_column("content")
331///     .query_vector(Embedding::new(vec![0.1, 0.2, 0.3]))
332///     .query_text("machine learning")
333///     .metric(DistanceMetric::Cosine)
334///     .vector_weight(0.7)
335///     .text_weight(0.3)
336///     .limit(10)
337///     .build();
338///
339/// let sql = query.to_sql();
340/// ```
341pub struct HybridSearchBuilder {
342    table: String,
343    vector_column: Option<String>,
344    text_column: Option<String>,
345    query_vector: Option<Embedding>,
346    query_text: Option<String>,
347    metric: DistanceMetric,
348    vector_weight: f64,
349    text_weight: f64,
350    limit: usize,
351    language: String,
352    where_clauses: Vec<String>,
353}
354
355impl HybridSearchBuilder {
356    /// Create a new hybrid search builder.
357    pub fn new(table: impl Into<String>) -> Self {
358        Self {
359            table: table.into(),
360            vector_column: None,
361            text_column: None,
362            query_vector: None,
363            query_text: None,
364            metric: DistanceMetric::Cosine,
365            vector_weight: 0.5,
366            text_weight: 0.5,
367            limit: 10,
368            language: "english".to_string(),
369            where_clauses: Vec::new(),
370        }
371    }
372
373    /// Set the vector column name.
374    pub fn vector_column(mut self, column: impl Into<String>) -> Self {
375        self.vector_column = Some(column.into());
376        self
377    }
378
379    /// Set the text column name.
380    pub fn text_column(mut self, column: impl Into<String>) -> Self {
381        self.text_column = Some(column.into());
382        self
383    }
384
385    /// Set the query vector.
386    pub fn query_vector(mut self, embedding: Embedding) -> Self {
387        self.query_vector = Some(embedding);
388        self
389    }
390
391    /// Set the text query.
392    pub fn query_text(mut self, text: impl Into<String>) -> Self {
393        self.query_text = Some(text.into());
394        self
395    }
396
397    /// Set the vector distance metric.
398    pub fn metric(mut self, metric: DistanceMetric) -> Self {
399        self.metric = metric;
400        self
401    }
402
403    /// Set the weight for the vector similarity component (0.0 to 1.0).
404    pub fn vector_weight(mut self, weight: f64) -> Self {
405        self.vector_weight = weight;
406        self
407    }
408
409    /// Set the weight for the text relevance component (0.0 to 1.0).
410    pub fn text_weight(mut self, weight: f64) -> Self {
411        self.text_weight = weight;
412        self
413    }
414
415    /// Set the result limit.
416    pub fn limit(mut self, limit: usize) -> Self {
417        self.limit = limit;
418        self
419    }
420
421    /// Set the text search language.
422    ///
423    /// The language is interpolated into SQL string literals, so it must be a
424    /// valid PostgreSQL text search configuration name: a simple identifier
425    /// matching `^[A-Za-z_][A-Za-z0-9_]*$` (e.g. `english`, `simple`).
426    ///
427    /// This is the panicking convenience for trusted, hardcoded values;
428    /// prefer [`try_language`](Self::try_language) when the language comes
429    /// from untrusted (e.g. request-influenced) input.
430    ///
431    /// # Panics
432    ///
433    /// Panics if `language` is not a valid configuration name.
434    pub fn language(mut self, language: impl Into<String>) -> Self {
435        let language = language.into();
436        assert!(
437            is_valid_ts_config_name(&language),
438            "invalid text search language {language:?}: must match ^[A-Za-z_][A-Za-z0-9_]*$"
439        );
440        self.language = language;
441        self
442    }
443
444    /// Set the text search language, returning an error for invalid names.
445    ///
446    /// Fallible counterpart to [`language`](Self::language) — use this when
447    /// the language is influenced by untrusted input so an invalid name is
448    /// a recoverable error instead of a panic.
449    pub fn try_language(mut self, language: impl Into<String>) -> VectorResult<Self> {
450        let language = language.into();
451        if !is_valid_ts_config_name(&language) {
452            return Err(VectorError::config(format!(
453                "invalid text search language {language:?}: must match ^[A-Za-z_][A-Za-z0-9_]*$"
454            )));
455        }
456        self.language = language;
457        Ok(self)
458    }
459
460    /// Add a WHERE condition.
461    pub fn where_clause(mut self, condition: impl Into<String>) -> Self {
462        self.where_clauses.push(condition.into());
463        self
464    }
465
466    /// Build the hybrid search query.
467    pub fn build(self) -> HybridSearchQuery {
468        HybridSearchQuery {
469            table: self.table,
470            vector_column: self
471                .vector_column
472                .unwrap_or_else(|| "embedding".to_string()),
473            text_column: self.text_column.unwrap_or_else(|| "content".to_string()),
474            query_vector: self.query_vector,
475            query_text: self.query_text,
476            metric: self.metric,
477            vector_weight: self.vector_weight,
478            text_weight: self.text_weight,
479            limit: self.limit,
480            language: self.language,
481            where_clauses: self.where_clauses,
482        }
483    }
484}
485
486/// A hybrid search query combining vector similarity and full-text search.
487///
488/// **Result shape:** the SQL from [`to_sql`](Self::to_sql) returns two columns —
489/// an anonymous composite named `coalesce` (the merged matched row) and
490/// `rrf_score`. The composite's fields cannot be expanded or selected in SQL;
491/// see [`to_sql`](Self::to_sql) for how to consume it.
492#[derive(Debug, Clone, Serialize, Deserialize)]
493pub struct HybridSearchQuery {
494    /// Table name.
495    pub table: String,
496    /// Vector column.
497    pub vector_column: String,
498    /// Text column.
499    pub text_column: String,
500    /// Query vector.
501    pub query_vector: Option<Embedding>,
502    /// Text query.
503    pub query_text: Option<String>,
504    /// Distance metric.
505    pub metric: DistanceMetric,
506    /// Weight for vector similarity (0.0-1.0).
507    pub vector_weight: f64,
508    /// Weight for text relevance (0.0-1.0).
509    pub text_weight: f64,
510    /// Result limit.
511    pub limit: usize,
512    /// Text search language.
513    pub language: String,
514    /// Additional WHERE conditions.
515    pub where_clauses: Vec<String>,
516}
517
518impl HybridSearchQuery {
519    /// Generate the SQL query using Reciprocal Rank Fusion (RRF).
520    ///
521    /// RRF combines rankings from multiple retrieval methods:
522    /// `score = sum(1 / (k + rank_i))` where k is a constant (typically 60).
523    ///
524    /// The query vector should be `$1` and the text query should be `$2`.
525    ///
526    /// # Result shape
527    ///
528    /// The final SELECT returns **two columns**: an anonymous composite named
529    /// `coalesce` — the whole matched row from `vector_results` (when both
530    /// sides match, the vector side wins) or `text_results` — plus `rrf_score`.
531    /// The composite's record type is never registered, so PostgreSQL rejects
532    /// both `(coalesce).*` expansion and per-field access like `(coalesce).id`
533    /// with "record type has not been registered". Consumers must either
534    /// decode the composite value client-side or wrap the query and project
535    /// `row_to_json(coalesce)` (or `to_json(coalesce)`). The composite also
536    /// carries the matching side's rank column (`vec_rank` or `text_rank`) in
537    /// addition to the base table columns, so its shape is not fixed.
538    pub fn to_sql(&self) -> String {
539        // The builder validates `language`, but these fields are public, so
540        // escape single quotes as defense-in-depth before interpolating into
541        // SQL string literals.
542        let lang = self.language.replace('\'', "''");
543        let vec_distance = format!("{} {} $1", self.vector_column, self.metric.operator());
544        let text_rank = format!(
545            "ts_rank(to_tsvector('{lang}', {}), plainto_tsquery('{lang}', $2))",
546            self.text_column
547        );
548        let fts_match = format!(
549            "to_tsvector('{lang}', {}) @@ plainto_tsquery('{lang}', $2)",
550            self.text_column
551        );
552
553        let where_clause = if self.where_clauses.is_empty() {
554            String::new()
555        } else {
556            format!(" WHERE {}", self.where_clauses.join(" AND "))
557        };
558
559        // text_results already filters on the FTS match, so user filters must
560        // be AND-joined to it — a second WHERE keyword would be invalid SQL.
561        let text_where = if self.where_clauses.is_empty() {
562            format!(" WHERE {fts_match}")
563        } else {
564            format!("{where_clause} AND {fts_match}")
565        };
566
567        // Use RRF scoring: combine vector and text rankings
568        format!(
569            "WITH vector_results AS (\
570                SELECT *, ROW_NUMBER() OVER (ORDER BY {vec_distance}) AS vec_rank \
571                FROM {table}{where_clause} \
572                ORDER BY {vec_distance} \
573                LIMIT {fetch_limit}\
574            ), \
575            text_results AS (\
576                SELECT *, ROW_NUMBER() OVER (ORDER BY {text_rank} DESC) AS text_rank \
577                FROM {table}{text_where} \
578                ORDER BY {text_rank} DESC \
579                LIMIT {fetch_limit}\
580            ) \
581            SELECT COALESCE(v.*, t.*), \
582                ({vec_weight} / (60.0 + COALESCE(v.vec_rank, 1000))) + \
583                ({text_weight} / (60.0 + COALESCE(t.text_rank, 1000))) AS rrf_score \
584            FROM vector_results v \
585            FULL OUTER JOIN text_results t ON v.id = t.id \
586            ORDER BY rrf_score DESC \
587            LIMIT {limit}",
588            table = self.table,
589            where_clause = where_clause,
590            text_where = text_where,
591            fetch_limit = self.limit * 3, // Fetch more for fusion
592            vec_weight = self.vector_weight,
593            text_weight = self.text_weight,
594            limit = self.limit,
595        )
596    }
597}
598
599#[cfg(test)]
600mod tests {
601    use super::*;
602
603    fn test_embedding() -> Embedding {
604        Embedding::new(vec![0.1, 0.2, 0.3])
605    }
606
607    #[test]
608    fn test_basic_search_query() {
609        let query = VectorSearchBuilder::new("documents", "embedding")
610            .query(test_embedding())
611            .metric(DistanceMetric::Cosine)
612            .limit(10)
613            .build();
614
615        let sql = query.to_sql();
616        assert!(sql.contains("SELECT *"));
617        assert!(sql.contains("AS distance"));
618        assert!(sql.contains("<=>"));
619        assert!(sql.contains("$1"));
620        assert!(sql.contains("FROM documents"));
621        assert!(sql.contains("LIMIT 10"));
622    }
623
624    #[test]
625    fn test_search_with_select() {
626        let query = VectorSearchBuilder::new("documents", "embedding")
627            .query(test_embedding())
628            .select(&["id", "title"])
629            .build();
630
631        let sql = query.to_sql();
632        assert!(sql.contains("SELECT id, title"));
633    }
634
635    #[test]
636    fn test_search_with_where() {
637        let query = VectorSearchBuilder::new("documents", "embedding")
638            .query(test_embedding())
639            .where_clause("category = 'tech'")
640            .where_clause("published = true")
641            .build();
642
643        let sql = query.to_sql();
644        assert!(sql.contains("WHERE"));
645        assert!(sql.contains("category = 'tech'"));
646        assert!(sql.contains("published = true"));
647        assert!(sql.contains("AND"));
648    }
649
650    #[test]
651    fn test_search_with_max_distance() {
652        let query = VectorSearchBuilder::new("documents", "embedding")
653            .query(test_embedding())
654            .metric(DistanceMetric::L2)
655            .max_distance(0.5)
656            .build();
657
658        let sql = query.to_sql();
659        assert!(sql.contains("< 0.5"));
660    }
661
662    #[test]
663    fn test_search_with_distance_range() {
664        let query = VectorSearchBuilder::new("documents", "embedding")
665            .query(test_embedding())
666            .min_distance(0.1)
667            .max_distance(0.5)
668            .build();
669
670        let sql = query.to_sql();
671        assert!(sql.contains("< 0.5"));
672        assert!(sql.contains(">= 0.1"));
673    }
674
675    #[test]
676    fn test_search_without_distance() {
677        let query = VectorSearchBuilder::new("documents", "embedding")
678            .query(test_embedding())
679            .without_distance()
680            .build();
681
682        let sql = query.to_sql();
683        assert!(!sql.contains("AS distance"));
684    }
685
686    #[test]
687    fn test_search_custom_alias() {
688        let query = VectorSearchBuilder::new("documents", "embedding")
689            .query(test_embedding())
690            .distance_alias("similarity")
691            .build();
692
693        let sql = query.to_sql();
694        assert!(sql.contains("AS similarity"));
695    }
696
697    #[test]
698    fn test_search_with_pagination() {
699        let query = VectorSearchBuilder::new("documents", "embedding")
700            .query(test_embedding())
701            .limit(10)
702            .offset(20)
703            .build();
704
705        let sql = query.to_sql();
706        assert!(sql.contains("LIMIT 10"));
707        assert!(sql.contains("OFFSET 20"));
708    }
709
710    #[test]
711    fn test_search_with_extra_order_by() {
712        let query = VectorSearchBuilder::new("documents", "embedding")
713            .query(test_embedding())
714            .then_order_by("created_at DESC")
715            .build();
716
717        let sql = query.to_sql();
718        assert!(sql.contains("ORDER BY distance, created_at DESC"));
719    }
720
721    #[test]
722    fn test_search_params() {
723        let query = VectorSearchBuilder::new("documents", "embedding")
724            .query(test_embedding())
725            .probes(10)
726            .ef_search(200)
727            .build();
728
729        let set_sql = query.param_set_sql();
730        assert_eq!(set_sql.len(), 2);
731        assert!(set_sql[0].contains("ivfflat.probes = 10"));
732        assert!(set_sql[1].contains("hnsw.ef_search = 200"));
733    }
734
735    #[test]
736    fn test_try_build_without_vector() {
737        let result = VectorSearchBuilder::new("documents", "embedding").try_build();
738        assert!(result.is_none());
739    }
740
741    #[test]
742    fn test_custom_param_index() {
743        let query = VectorSearchBuilder::new("documents", "embedding")
744            .query(test_embedding())
745            .build();
746
747        let sql = query.to_sql_with_param(3);
748        assert!(sql.contains("$3"));
749    }
750
751    #[test]
752    fn test_hybrid_search() {
753        let query = HybridSearchBuilder::new("documents")
754            .vector_column("embedding")
755            .text_column("content")
756            .query_vector(test_embedding())
757            .query_text("machine learning")
758            .metric(DistanceMetric::Cosine)
759            .vector_weight(0.7)
760            .text_weight(0.3)
761            .limit(10)
762            .build();
763
764        let sql = query.to_sql();
765        assert!(sql.contains("vector_results"));
766        assert!(sql.contains("text_results"));
767        assert!(sql.contains("rrf_score"));
768        assert!(sql.contains("<=>"));
769        assert!(sql.contains("ts_rank"));
770        assert!(sql.contains("FULL OUTER JOIN"));
771    }
772
773    #[test]
774    fn test_hybrid_search_with_where_clause_single_where() {
775        let query = HybridSearchBuilder::new("documents")
776            .vector_column("embedding")
777            .text_column("content")
778            .query_vector(test_embedding())
779            .query_text("machine learning")
780            .where_clause("category = 'tech'")
781            .build();
782
783        let sql = query.to_sql();
784        // One WHERE per CTE: the FTS predicate must be AND-joined to the user
785        // filter in text_results, not emitted as a second WHERE keyword.
786        assert_eq!(sql.matches("WHERE").count(), 2);
787        assert!(sql.contains(
788            "WHERE category = 'tech' AND to_tsvector('english', content) @@ plainto_tsquery('english', $2)"
789        ));
790    }
791
792    #[test]
793    #[should_panic(expected = "invalid text search language")]
794    fn test_hybrid_language_rejects_invalid() {
795        let _ = HybridSearchBuilder::new("documents").language("english'); DROP TABLE users; --");
796    }
797
798    #[test]
799    fn test_hybrid_try_language_returns_error_instead_of_panicking() {
800        let result =
801            HybridSearchBuilder::new("documents").try_language("english'); DROP TABLE users; --");
802        assert!(result.is_err(), "invalid language should be an error");
803
804        let builder = HybridSearchBuilder::new("documents")
805            .try_language("spanish")
806            .expect("valid language should build");
807        assert_eq!(builder.build().language, "spanish");
808    }
809
810    #[test]
811    fn test_hybrid_language_escaped_when_constructed_directly() {
812        // HybridSearchQuery's fields are public, so to_sql must still escape
813        // single quotes as defense-in-depth.
814        let mut query = HybridSearchBuilder::new("documents").build();
815        query.language = "english'); DROP TABLE users; --".to_string();
816        let sql = query.to_sql();
817        assert!(sql.contains("'english''); DROP TABLE users; --'"));
818        assert!(!sql.contains("'english');"));
819    }
820
821    #[test]
822    fn test_all_metrics_produce_valid_sql() {
823        for metric in [
824            DistanceMetric::L2,
825            DistanceMetric::InnerProduct,
826            DistanceMetric::Cosine,
827            DistanceMetric::L1,
828        ] {
829            let query = VectorSearchBuilder::new("t", "c")
830                .query(test_embedding())
831                .metric(metric)
832                .build();
833            let sql = query.to_sql();
834            assert!(sql.contains(metric.operator()), "failed for {metric}");
835        }
836    }
837}
prax_pgvector/query.rs

prax_pgvector/
query.rs
