Skip to main content

dk_engine/graph/
symbols.rs

1use std::path::PathBuf;
2
3use dk_core::{RepoId, Span, Symbol, SymbolId, SymbolKind, Visibility};
4use sqlx::postgres::PgPool;
5use uuid::Uuid;
6
7/// Intermediate row type for mapping between database rows and `Symbol`.
8#[derive(sqlx::FromRow)]
9#[allow(dead_code)]
10struct SymbolRow {
11    id: Uuid,
12    repo_id: Uuid,
13    name: String,
14    qualified_name: String,
15    kind: String,
16    visibility: String,
17    file_path: String,
18    start_byte: i32,
19    end_byte: i32,
20    signature: Option<String>,
21    doc_comment: Option<String>,
22    parent_id: Option<Uuid>,
23    last_modified_by: Option<String>,
24    last_modified_intent: Option<String>,
25}
26
27impl SymbolRow {
28    fn into_symbol(self) -> Symbol {
29        Symbol {
30            id: self.id,
31            name: self.name,
32            qualified_name: self.qualified_name,
33            kind: self.kind.parse::<SymbolKind>().unwrap_or_else(|e| {
34                tracing::warn!("{e}, defaulting to Variable");
35                SymbolKind::Variable
36            }),
37            visibility: self.visibility.parse::<Visibility>().unwrap_or_else(|e| {
38                tracing::warn!("{e}, defaulting to Private");
39                Visibility::Private
40            }),
41            file_path: PathBuf::from(self.file_path),
42            span: Span {
43                start_byte: self.start_byte as u32,
44                end_byte: self.end_byte as u32,
45            },
46            signature: self.signature,
47            doc_comment: self.doc_comment,
48            parent: self.parent_id,
49            last_modified_by: self.last_modified_by,
50            last_modified_intent: self.last_modified_intent,
51        }
52    }
53}
54
55/// PostgreSQL-backed CRUD store for the symbol table.
56#[derive(Clone)]
57pub struct SymbolStore {
58    pool: PgPool,
59}
60
61impl SymbolStore {
62    /// Create a new `SymbolStore` backed by the given connection pool.
63    pub fn new(pool: PgPool) -> Self {
64        Self { pool }
65    }
66
67    /// Insert or update a symbol.
68    ///
69    /// Uses `ON CONFLICT (repo_id, qualified_name) DO UPDATE` so that
70    /// repeated ingestion of the same file is idempotent.
71    pub async fn upsert_symbol(
72        &self,
73        repo_id: RepoId,
74        sym: &Symbol,
75    ) -> dk_core::Result<()> {
76        let kind_str = sym.kind.to_string();
77        let vis_str = sym.visibility.to_string();
78        let file_path_str = sym.file_path.to_string_lossy().to_string();
79
80        sqlx::query(
81            r#"
82            INSERT INTO symbols (
83                id, repo_id, name, qualified_name, kind, visibility,
84                file_path, start_byte, end_byte, signature, doc_comment,
85                parent_id, last_modified_by, last_modified_intent
86            )
87            VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)
88            ON CONFLICT (repo_id, qualified_name) DO UPDATE SET
89                id = EXCLUDED.id,
90                name = EXCLUDED.name,
91                kind = EXCLUDED.kind,
92                visibility = EXCLUDED.visibility,
93                file_path = EXCLUDED.file_path,
94                start_byte = EXCLUDED.start_byte,
95                end_byte = EXCLUDED.end_byte,
96                signature = EXCLUDED.signature,
97                doc_comment = EXCLUDED.doc_comment,
98                parent_id = EXCLUDED.parent_id,
99                last_modified_by = EXCLUDED.last_modified_by,
100                last_modified_intent = EXCLUDED.last_modified_intent
101            "#,
102        )
103        .bind(sym.id)
104        .bind(repo_id)
105        .bind(&sym.name)
106        .bind(&sym.qualified_name)
107        .bind(&kind_str)
108        .bind(&vis_str)
109        .bind(&file_path_str)
110        .bind(sym.span.start_byte as i32)
111        .bind(sym.span.end_byte as i32)
112        .bind(&sym.signature)
113        .bind(&sym.doc_comment)
114        .bind(sym.parent)
115        .bind(&sym.last_modified_by)
116        .bind(&sym.last_modified_intent)
117        .execute(&self.pool)
118        .await?;
119
120        Ok(())
121    }
122
123    /// Search symbols by name or qualified_name using ILIKE.
124    pub async fn find_symbols(
125        &self,
126        repo_id: RepoId,
127        query: &str,
128    ) -> dk_core::Result<Vec<Symbol>> {
129        let pattern = format!("%{query}%");
130        let rows = sqlx::query_as::<_, SymbolRow>(
131            r#"
132            SELECT id, repo_id, name, qualified_name, kind, visibility,
133                   file_path, start_byte, end_byte, signature, doc_comment,
134                   parent_id, last_modified_by, last_modified_intent
135            FROM symbols
136            WHERE repo_id = $1 AND (name ILIKE $2 OR qualified_name ILIKE $2)
137            ORDER BY qualified_name
138            "#,
139        )
140        .bind(repo_id)
141        .bind(&pattern)
142        .fetch_all(&self.pool)
143        .await?;
144
145        Ok(rows.into_iter().map(SymbolRow::into_symbol).collect())
146    }
147
148    /// Find all symbols of a given kind in a repository.
149    pub async fn find_by_kind(
150        &self,
151        repo_id: RepoId,
152        kind: &SymbolKind,
153    ) -> dk_core::Result<Vec<Symbol>> {
154        let kind_str = kind.to_string();
155        let rows = sqlx::query_as::<_, SymbolRow>(
156            r#"
157            SELECT id, repo_id, name, qualified_name, kind, visibility,
158                   file_path, start_byte, end_byte, signature, doc_comment,
159                   parent_id, last_modified_by, last_modified_intent
160            FROM symbols
161            WHERE repo_id = $1 AND kind = $2
162            ORDER BY qualified_name
163            "#,
164        )
165        .bind(repo_id)
166        .bind(&kind_str)
167        .fetch_all(&self.pool)
168        .await?;
169
170        Ok(rows.into_iter().map(SymbolRow::into_symbol).collect())
171    }
172
173    /// Find all symbols in a given file.
174    pub async fn find_by_file(
175        &self,
176        repo_id: RepoId,
177        file_path: &str,
178    ) -> dk_core::Result<Vec<Symbol>> {
179        let rows = sqlx::query_as::<_, SymbolRow>(
180            r#"
181            SELECT id, repo_id, name, qualified_name, kind, visibility,
182                   file_path, start_byte, end_byte, signature, doc_comment,
183                   parent_id, last_modified_by, last_modified_intent
184            FROM symbols
185            WHERE repo_id = $1 AND file_path = $2
186            ORDER BY start_byte
187            "#,
188        )
189        .bind(repo_id)
190        .bind(file_path)
191        .fetch_all(&self.pool)
192        .await?;
193
194        Ok(rows.into_iter().map(SymbolRow::into_symbol).collect())
195    }
196
197    /// Get a single symbol by its ID.
198    pub async fn get_by_id(&self, id: SymbolId) -> dk_core::Result<Option<Symbol>> {
199        let row = sqlx::query_as::<_, SymbolRow>(
200            r#"
201            SELECT id, repo_id, name, qualified_name, kind, visibility,
202                   file_path, start_byte, end_byte, signature, doc_comment,
203                   parent_id, last_modified_by, last_modified_intent
204            FROM symbols
205            WHERE id = $1
206            "#,
207        )
208        .bind(id)
209        .fetch_optional(&self.pool)
210        .await?;
211
212        Ok(row.map(SymbolRow::into_symbol))
213    }
214
215    /// Fetch multiple symbols by their IDs in a single batch query.
216    ///
217    /// Returns symbols in arbitrary order. Symbols that do not exist are
218    /// silently omitted.
219    pub async fn get_by_ids(&self, ids: &[SymbolId]) -> dk_core::Result<Vec<Symbol>> {
220        if ids.is_empty() {
221            return Ok(Vec::new());
222        }
223
224        let rows = sqlx::query_as::<_, SymbolRow>(
225            r#"
226            SELECT id, repo_id, name, qualified_name, kind, visibility,
227                   file_path, start_byte, end_byte, signature, doc_comment,
228                   parent_id, last_modified_by, last_modified_intent
229            FROM symbols
230            WHERE id = ANY($1)
231            "#,
232        )
233        .bind(ids)
234        .fetch_all(&self.pool)
235        .await?;
236
237        Ok(rows.into_iter().map(SymbolRow::into_symbol).collect())
238    }
239
240    /// Delete all symbols belonging to a file. Returns the number of rows deleted.
241    pub async fn delete_by_file(
242        &self,
243        repo_id: RepoId,
244        file_path: &str,
245    ) -> dk_core::Result<u64> {
246        let result = sqlx::query(
247            "DELETE FROM symbols WHERE repo_id = $1 AND file_path = $2",
248        )
249        .bind(repo_id)
250        .bind(file_path)
251        .execute(&self.pool)
252        .await?;
253
254        Ok(result.rows_affected())
255    }
256
257    /// Delete all symbols belonging to a repository. Returns the number of rows deleted.
258    pub async fn delete_by_repo(&self, repo_id: RepoId) -> dk_core::Result<u64> {
259        let result = sqlx::query("DELETE FROM symbols WHERE repo_id = $1")
260            .bind(repo_id)
261            .execute(&self.pool)
262            .await?;
263
264        Ok(result.rows_affected())
265    }
266
267    /// Count symbols in a repository.
268    pub async fn count(&self, repo_id: RepoId) -> dk_core::Result<i64> {
269        let (count,): (i64,) =
270            sqlx::query_as("SELECT COUNT(*) FROM symbols WHERE repo_id = $1")
271                .bind(repo_id)
272                .fetch_one(&self.pool)
273                .await?;
274
275        Ok(count)
276    }
277}