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 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 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 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 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 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 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 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 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 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}