1use std::path::PathBuf;
2
3use dk_core::{RepoId, Span, Symbol, SymbolId, SymbolKind, Visibility};
4use sqlx::postgres::PgPool;
5use uuid::Uuid;
6
7#[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#[derive(Clone)]
57pub struct SymbolStore {
58 pool: PgPool,
59}
60
61impl SymbolStore {
62 pub fn new(pool: PgPool) -> Self {
64 Self { pool }
65 }
66
67 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 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 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 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 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 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 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 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 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}