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                name = EXCLUDED.name,
90                kind = EXCLUDED.kind,
91                visibility = EXCLUDED.visibility,
92                file_path = EXCLUDED.file_path,
93                start_byte = EXCLUDED.start_byte,
94                end_byte = EXCLUDED.end_byte,
95                signature = EXCLUDED.signature,
96                doc_comment = EXCLUDED.doc_comment,
97                parent_id = EXCLUDED.parent_id,
98                last_modified_by = EXCLUDED.last_modified_by,
99                last_modified_intent = EXCLUDED.last_modified_intent
100            "#,
101        )
102        .bind(sym.id)
103        .bind(repo_id)
104        .bind(&sym.name)
105        .bind(&sym.qualified_name)
106        .bind(&kind_str)
107        .bind(&vis_str)
108        .bind(&file_path_str)
109        .bind(sym.span.start_byte as i32)
110        .bind(sym.span.end_byte as i32)
111        .bind(&sym.signature)
112        .bind(&sym.doc_comment)
113        .bind(sym.parent)
114        .bind(&sym.last_modified_by)
115        .bind(&sym.last_modified_intent)
116        .execute(&self.pool)
117        .await?;
118
119        Ok(())
120    }
121
122    /// Search symbols by name or qualified_name using ILIKE.
123    pub async fn find_symbols(
124        &self,
125        repo_id: RepoId,
126        query: &str,
127    ) -> dk_core::Result<Vec<Symbol>> {
128        let pattern = format!("%{query}%");
129        let rows = sqlx::query_as::<_, SymbolRow>(
130            r#"
131            SELECT id, repo_id, name, qualified_name, kind, visibility,
132                   file_path, start_byte, end_byte, signature, doc_comment,
133                   parent_id, last_modified_by, last_modified_intent
134            FROM symbols
135            WHERE repo_id = $1 AND (name ILIKE $2 OR qualified_name ILIKE $2)
136            ORDER BY qualified_name
137            "#,
138        )
139        .bind(repo_id)
140        .bind(&pattern)
141        .fetch_all(&self.pool)
142        .await?;
143
144        Ok(rows.into_iter().map(SymbolRow::into_symbol).collect())
145    }
146
147    /// Find all symbols of a given kind in a repository.
148    pub async fn find_by_kind(
149        &self,
150        repo_id: RepoId,
151        kind: &SymbolKind,
152    ) -> dk_core::Result<Vec<Symbol>> {
153        let kind_str = kind.to_string();
154        let rows = sqlx::query_as::<_, SymbolRow>(
155            r#"
156            SELECT id, repo_id, name, qualified_name, kind, visibility,
157                   file_path, start_byte, end_byte, signature, doc_comment,
158                   parent_id, last_modified_by, last_modified_intent
159            FROM symbols
160            WHERE repo_id = $1 AND kind = $2
161            ORDER BY qualified_name
162            "#,
163        )
164        .bind(repo_id)
165        .bind(&kind_str)
166        .fetch_all(&self.pool)
167        .await?;
168
169        Ok(rows.into_iter().map(SymbolRow::into_symbol).collect())
170    }
171
172    /// Find all symbols in a given file.
173    pub async fn find_by_file(
174        &self,
175        repo_id: RepoId,
176        file_path: &str,
177    ) -> dk_core::Result<Vec<Symbol>> {
178        let rows = sqlx::query_as::<_, SymbolRow>(
179            r#"
180            SELECT id, repo_id, name, qualified_name, kind, visibility,
181                   file_path, start_byte, end_byte, signature, doc_comment,
182                   parent_id, last_modified_by, last_modified_intent
183            FROM symbols
184            WHERE repo_id = $1 AND file_path = $2
185            ORDER BY start_byte
186            "#,
187        )
188        .bind(repo_id)
189        .bind(file_path)
190        .fetch_all(&self.pool)
191        .await?;
192
193        Ok(rows.into_iter().map(SymbolRow::into_symbol).collect())
194    }
195
196    /// Get a single symbol by its ID.
197    pub async fn get_by_id(&self, id: SymbolId) -> dk_core::Result<Option<Symbol>> {
198        let row = sqlx::query_as::<_, SymbolRow>(
199            r#"
200            SELECT id, repo_id, name, qualified_name, kind, visibility,
201                   file_path, start_byte, end_byte, signature, doc_comment,
202                   parent_id, last_modified_by, last_modified_intent
203            FROM symbols
204            WHERE id = $1
205            "#,
206        )
207        .bind(id)
208        .fetch_optional(&self.pool)
209        .await?;
210
211        Ok(row.map(SymbolRow::into_symbol))
212    }
213
214    /// Fetch multiple symbols by their IDs in a single batch query.
215    ///
216    /// Returns symbols in arbitrary order. Symbols that do not exist are
217    /// silently omitted.
218    pub async fn get_by_ids(&self, ids: &[SymbolId]) -> dk_core::Result<Vec<Symbol>> {
219        if ids.is_empty() {
220            return Ok(Vec::new());
221        }
222
223        let rows = sqlx::query_as::<_, SymbolRow>(
224            r#"
225            SELECT id, repo_id, name, qualified_name, kind, visibility,
226                   file_path, start_byte, end_byte, signature, doc_comment,
227                   parent_id, last_modified_by, last_modified_intent
228            FROM symbols
229            WHERE id = ANY($1)
230            "#,
231        )
232        .bind(ids)
233        .fetch_all(&self.pool)
234        .await?;
235
236        Ok(rows.into_iter().map(SymbolRow::into_symbol).collect())
237    }
238
239    /// Delete all symbols belonging to a file. Returns the number of rows deleted.
240    pub async fn delete_by_file(
241        &self,
242        repo_id: RepoId,
243        file_path: &str,
244    ) -> dk_core::Result<u64> {
245        let result = sqlx::query(
246            "DELETE FROM symbols WHERE repo_id = $1 AND file_path = $2",
247        )
248        .bind(repo_id)
249        .bind(file_path)
250        .execute(&self.pool)
251        .await?;
252
253        Ok(result.rows_affected())
254    }
255
256    /// Delete all symbols belonging to a repository. Returns the number of rows deleted.
257    pub async fn delete_by_repo(&self, repo_id: RepoId) -> dk_core::Result<u64> {
258        let result = sqlx::query("DELETE FROM symbols WHERE repo_id = $1")
259            .bind(repo_id)
260            .execute(&self.pool)
261            .await?;
262
263        Ok(result.rows_affected())
264    }
265
266    /// Count symbols in a repository.
267    pub async fn count(&self, repo_id: RepoId) -> dk_core::Result<i64> {
268        let (count,): (i64,) =
269            sqlx::query_as("SELECT COUNT(*) FROM symbols WHERE repo_id = $1")
270                .bind(repo_id)
271                .fetch_one(&self.pool)
272                .await?;
273
274        Ok(count)
275    }
276}