sbol_db_postgres/repo/
sequence_search.rs

1//! Nucleotide substring + reverse-complement search over `sbol_sequences`.
2//!
3//! The seed index is `sbol_sequence_kmers`: one row per observed canonical
4//! 8-mer position. At query time we pick the first k-mer of the query and
5//! of its reverse complement, ask Postgres for the candidate sequence ids,
6//! then verify each candidate by direct substring match on
7//! `sbol_sequences.elements` in Rust. Queries shorter than K fall back to
8//! an `ILIKE` candidate scan (no seed possible at that length).
9//!
10//! Indexing is invoked from [`crate::repo::projections::upsert_sequence`]
11//! during import so the seed index stays in lockstep with the typed
12//! projection.
13
14use sbol_db_core::{
15    kmer::{
16        canonical, canonical_kmers, encode_kmer, reverse_complement_string, KmerStrand, KMER_K,
17    },
18    DomainError,
19};
20use sqlx::{PgConnection, Row};
21use uuid::Uuid;
22
23use crate::repo::db_err;
24use crate::PgPool;
25
26use sbol_db_storage::{BatchSequenceMatch, SequenceMatch, SequenceSearchOptions};
27
28#[derive(Clone)]
29pub struct SequenceSearchRepository {
30    pool: PgPool,
31}
32
33impl SequenceSearchRepository {
34    pub fn new(pool: PgPool) -> Self {
35        Self { pool }
36    }
37
38    /// Rebuild the k-mer index for one sequence. Idempotent: deletes all
39    /// existing rows then inserts the current set. Skips sequences whose
40    /// `alphabet` is not DNA/RNA, whose `elements` is empty, or whose length
41    /// is shorter than `K`.
42    pub async fn reindex(
43        &self,
44        conn: &mut PgConnection,
45        object_id: Uuid,
46        elements: Option<&str>,
47        alphabet: Option<&str>,
48    ) -> Result<usize, DomainError> {
49        reindex_kmers(conn, object_id, elements, alphabet).await
50    }
51
52    /// Exact substring search with reverse-complement awareness. Returns
53    /// every (sequence, start, strand) where the pattern (or its reverse
54    /// complement) appears in `sbol_sequences.elements`.
55    pub async fn search(
56        &self,
57        pattern: &str,
58        options: SequenceSearchOptions,
59    ) -> Result<Vec<SequenceMatch>, DomainError> {
60        let max = options.max_hits.unwrap_or(1024) as usize;
61        if max == 0 {
62            return Ok(Vec::new());
63        }
64        let q_upper: String = pattern
65            .chars()
66            .filter(|c| !c.is_whitespace())
67            .map(|c| c.to_ascii_uppercase())
68            .collect();
69        if q_upper.is_empty() {
70            return Ok(Vec::new());
71        }
72        let q_rc: String = reverse_complement_string(&q_upper).to_ascii_uppercase();
73        let include_rc = !matches!(options.forward_only, Some(true));
74        let candidates = if q_upper.len() < KMER_K {
75            self.candidates_short(&q_upper, &q_rc, include_rc).await?
76        } else {
77            self.candidates_seeded(&q_upper, &q_rc, include_rc).await?
78        };
79
80        let mut matches = Vec::new();
81        'outer: for (iri, elements) in candidates {
82            let up = elements.to_ascii_uppercase();
83            for (start, _) in up.match_indices(q_upper.as_str()) {
84                matches.push(SequenceMatch {
85                    sequence_iri: iri.clone(),
86                    start: start as i32,
87                    length: q_upper.len() as i32,
88                    strand: KmerStrand::Forward.as_char(),
89                });
90                if matches.len() >= max {
91                    break 'outer;
92                }
93            }
94            if include_rc && q_upper != q_rc {
95                for (start, _) in up.match_indices(q_rc.as_str()) {
96                    matches.push(SequenceMatch {
97                        sequence_iri: iri.clone(),
98                        start: start as i32,
99                        length: q_rc.len() as i32,
100                        strand: KmerStrand::Reverse.as_char(),
101                    });
102                    if matches.len() >= max {
103                        break 'outer;
104                    }
105                }
106            }
107        }
108        Ok(matches)
109    }
110
111    /// Run [`search`](Self::search) for each pattern. Patterns are independent;
112    /// the implementation is a simple loop. Use this to amortise HTTP / CLI
113    /// round-trips when scanning N motifs against the same corpus.
114    pub async fn search_many(
115        &self,
116        patterns: &[String],
117        options: SequenceSearchOptions,
118    ) -> Result<Vec<BatchSequenceMatch>, DomainError> {
119        let mut out = Vec::with_capacity(patterns.len());
120        for pattern in patterns {
121            let matches = self.search(pattern, options.clone()).await?;
122            out.push(BatchSequenceMatch {
123                pattern: pattern.clone(),
124                matches,
125            });
126        }
127        Ok(out)
128    }
129
130    /// Candidate sequences for the banded aligner: every indexed DNA/RNA
131    /// sequence sharing at least one canonical k-mer with `query`. Unlike
132    /// [`candidates_seeded`](Self::candidates_seeded), which seeds only on the
133    /// query's first k-mer, this aggregates over every canonical query k-mer so
134    /// a gapped hit whose head does not match is still admitted. A query shorter
135    /// than `K` has no seed and every nucleotide sequence is a candidate.
136    pub async fn align_candidates(
137        &self,
138        query: &str,
139    ) -> Result<Vec<(String, String)>, DomainError> {
140        let normalised: String = query.chars().filter(|c| !c.is_whitespace()).collect();
141        if normalised.is_empty() {
142            return Ok(Vec::new());
143        }
144        if normalised.len() < KMER_K {
145            return self.all_nucleotide_sequences().await;
146        }
147        let mut kmers: Vec<i32> = canonical_kmers(&normalised)
148            .map(|h| h.canonical as i32)
149            .collect();
150        kmers.sort_unstable();
151        kmers.dedup();
152        if kmers.is_empty() {
153            return self.all_nucleotide_sequences().await;
154        }
155
156        let rows = sqlx::query(
157            r#"
158            SELECT s.iri::text AS iri, s.elements AS elements
159            FROM sbol_sequences s
160            WHERE s.elements IS NOT NULL
161              AND s.alphabet IN ('DNA', 'RNA')
162              AND s.object_id IN (
163                SELECT DISTINCT sequence_object_id FROM sbol_sequence_kmers
164                WHERE kmer = ANY($1::int[])
165              )
166            "#,
167        )
168        .bind(&kmers)
169        .fetch_all(&self.pool)
170        .await
171        .map_err(db_err)?;
172        collect_rows(rows)
173    }
174
175    pub async fn all_nucleotide_sequences(&self) -> Result<Vec<(String, String)>, DomainError> {
176        let rows = sqlx::query(
177            r#"
178            SELECT iri::text AS iri, elements
179            FROM sbol_sequences
180            WHERE elements IS NOT NULL
181              AND alphabet IN ('DNA', 'RNA')
182            "#,
183        )
184        .fetch_all(&self.pool)
185        .await
186        .map_err(db_err)?;
187        collect_rows(rows)
188    }
189
190    pub async fn sequences_by_iris(
191        &self,
192        iris: &[String],
193    ) -> Result<Vec<(String, String)>, DomainError> {
194        if iris.is_empty() {
195            return Ok(Vec::new());
196        }
197        let rows = sqlx::query(
198            r#"
199            SELECT iri::text AS iri, elements
200            FROM sbol_sequences
201            WHERE elements IS NOT NULL
202              AND alphabet IN ('DNA', 'RNA')
203              AND iri::text = ANY($1::text[])
204            "#,
205        )
206        .bind(iris)
207        .fetch_all(&self.pool)
208        .await
209        .map_err(db_err)?;
210        collect_rows(rows)
211    }
212
213    async fn candidates_seeded(
214        &self,
215        q_upper: &str,
216        q_rc: &str,
217        include_rc: bool,
218    ) -> Result<Vec<(String, String)>, DomainError> {
219        let mut seeds: Vec<i32> = Vec::new();
220        if let Some(packed) = encode_kmer(&q_upper[..KMER_K]) {
221            seeds.push(canonical(packed).0 as i32);
222        }
223        if include_rc && q_upper != q_rc {
224            if let Some(packed) = encode_kmer(&q_rc[..KMER_K]) {
225                let c = canonical(packed).0 as i32;
226                if !seeds.contains(&c) {
227                    seeds.push(c);
228                }
229            }
230        }
231        if seeds.is_empty() {
232            return Ok(Vec::new());
233        }
234        let rows = sqlx::query(
235            r#"
236            SELECT s.iri::text AS iri, s.elements AS elements
237            FROM sbol_sequences s
238            WHERE s.elements IS NOT NULL
239              AND s.alphabet IN ('DNA', 'RNA')
240              AND s.object_id IN (
241                SELECT DISTINCT sequence_object_id FROM sbol_sequence_kmers
242                WHERE kmer = ANY($1::int[])
243              )
244            "#,
245        )
246        .bind(&seeds)
247        .fetch_all(&self.pool)
248        .await
249        .map_err(db_err)?;
250
251        let mut out = Vec::with_capacity(rows.len());
252        for row in rows {
253            let iri: String = row.try_get("iri").map_err(db_err)?;
254            let elements: String = row.try_get("elements").map_err(db_err)?;
255            out.push((iri, elements));
256        }
257        Ok(out)
258    }
259
260    async fn candidates_short(
261        &self,
262        q_upper: &str,
263        q_rc: &str,
264        include_rc: bool,
265    ) -> Result<Vec<(String, String)>, DomainError> {
266        let sql = if include_rc && q_upper != q_rc {
267            r#"
268            SELECT iri::text AS iri, elements
269            FROM sbol_sequences
270            WHERE elements IS NOT NULL
271              AND alphabet IN ('DNA', 'RNA')
272              AND (UPPER(elements) LIKE '%' || $1 || '%'
273                   OR UPPER(elements) LIKE '%' || $2 || '%')
274            "#
275        } else {
276            r#"
277            SELECT iri::text AS iri, elements
278            FROM sbol_sequences
279            WHERE elements IS NOT NULL
280              AND alphabet IN ('DNA', 'RNA')
281              AND UPPER(elements) LIKE '%' || $1 || '%'
282            "#
283        };
284        let q = sqlx::query(sql).bind(q_upper);
285        let q = if include_rc && q_upper != q_rc {
286            q.bind(q_rc)
287        } else {
288            q
289        };
290        let rows = q.fetch_all(&self.pool).await.map_err(db_err)?;
291        let mut out = Vec::with_capacity(rows.len());
292        for row in rows {
293            let iri: String = row.try_get("iri").map_err(db_err)?;
294            let elements: String = row.try_get("elements").map_err(db_err)?;
295            out.push((iri, elements));
296        }
297        Ok(out)
298    }
299}
300
301/// Collect `(iri, elements)` from candidate rows selecting `iri` and
302/// `elements` columns.
303fn collect_rows(rows: Vec<sqlx::postgres::PgRow>) -> Result<Vec<(String, String)>, DomainError> {
304    let mut out = Vec::with_capacity(rows.len());
305    for row in rows {
306        let iri: String = row.try_get("iri").map_err(db_err)?;
307        let elements: String = row.try_get("elements").map_err(db_err)?;
308        out.push((iri, elements));
309    }
310    Ok(out)
311}
312
313/// Free-function form of the k-mer reindex, callable from
314/// `TypedProjectionRepository::upsert_sequence` without holding a repository
315/// instance. The caller passes the pooled connection used for the rest of
316/// the import transaction.
317pub(crate) async fn reindex_kmers(
318    conn: &mut PgConnection,
319    object_id: Uuid,
320    elements: Option<&str>,
321    alphabet: Option<&str>,
322) -> Result<usize, DomainError> {
323    sqlx::query("DELETE FROM sbol_sequence_kmers WHERE sequence_object_id = $1")
324        .bind(object_id)
325        .execute(&mut *conn)
326        .await
327        .map_err(db_err)?;
328    let Some(elements) = elements else {
329        return Ok(0);
330    };
331    if !matches!(alphabet, Some("DNA") | Some("RNA")) {
332        return Ok(0);
333    }
334    let normalised: String = elements.chars().filter(|c| !c.is_whitespace()).collect();
335    if normalised.len() < KMER_K {
336        return Ok(0);
337    }
338
339    let hits: Vec<_> = canonical_kmers(&normalised).collect();
340    if hits.is_empty() {
341        return Ok(0);
342    }
343
344    // Bulk insert via UNNEST -- one round trip even for long sequences.
345    let mut kmers: Vec<i32> = Vec::with_capacity(hits.len());
346    let mut positions: Vec<i32> = Vec::with_capacity(hits.len());
347    let mut strands: Vec<String> = Vec::with_capacity(hits.len());
348    for h in &hits {
349        kmers.push(h.canonical as i32);
350        positions.push(h.position as i32);
351        strands.push(h.strand.as_char().to_string());
352    }
353
354    sqlx::query(
355        r#"
356        INSERT INTO sbol_sequence_kmers (sequence_object_id, kmer, position, strand)
357        SELECT $1, k, p, s
358        FROM UNNEST($2::int[], $3::int[], $4::text[]) AS u(k, p, s)
359        "#,
360    )
361    .bind(object_id)
362    .bind(&kmers)
363    .bind(&positions)
364    .bind(&strands)
365    .execute(&mut *conn)
366    .await
367    .map_err(db_err)?;
368
369    Ok(hits.len())
370}
sbol_db_postgres/repo/sequence_search.rs

sbol_db_postgres/repo/
sequence_search.rs
