Skip to main content

cc_switch_lib/database/dao/
providers.rs

1//! 供应商数据访问对象
2//!
3//! 提供供应商(Provider)的 CRUD 操作。
4
5use crate::database::dao::providers_seed::{is_official_seed_id, OFFICIAL_SEEDS};
6use crate::database::{lock_conn, Database};
7use crate::error::AppError;
8use crate::provider::{Provider, ProviderMeta};
9use indexmap::IndexMap;
10use rusqlite::params;
11use std::collections::{HashMap, HashSet};
12
13impl Database {
14    /// 获取指定应用类型的所有供应商
15    pub fn get_all_providers(
16        &self,
17        app_type: &str,
18    ) -> Result<IndexMap<String, Provider>, AppError> {
19        let conn = lock_conn!(self.conn);
20        let mut stmt = conn.prepare(
21            "SELECT id, name, settings_config, website_url, category, created_at, sort_index, notes, icon, icon_color, meta, in_failover_queue
22             FROM providers WHERE app_type = ?1
23             ORDER BY COALESCE(sort_index, 999999), created_at ASC, id ASC"
24        ).map_err(|e| AppError::Database(e.to_string()))?;
25
26        let provider_iter = stmt
27            .query_map(params![app_type], |row| {
28                let id: String = row.get(0)?;
29                let name: String = row.get(1)?;
30                let settings_config_str: String = row.get(2)?;
31                let website_url: Option<String> = row.get(3)?;
32                let category: Option<String> = row.get(4)?;
33                let created_at: Option<i64> = row.get(5)?;
34                let sort_index: Option<usize> = row.get(6)?;
35                let notes: Option<String> = row.get(7)?;
36                let icon: Option<String> = row.get(8)?;
37                let icon_color: Option<String> = row.get(9)?;
38                let meta_str: String = row.get(10)?;
39                let in_failover_queue: bool = row.get(11)?;
40
41                let settings_config =
42                    serde_json::from_str(&settings_config_str).unwrap_or(serde_json::Value::Null);
43                let meta: ProviderMeta = serde_json::from_str(&meta_str).unwrap_or_default();
44
45                Ok((
46                    id,
47                    Provider {
48                        id: "".to_string(), // Placeholder, set below
49                        name,
50                        settings_config,
51                        website_url,
52                        category,
53                        created_at,
54                        sort_index,
55                        notes,
56                        meta: Some(meta),
57                        icon,
58                        icon_color,
59                        in_failover_queue,
60                    },
61                ))
62            })
63            .map_err(|e| AppError::Database(e.to_string()))?;
64
65        let mut providers = IndexMap::new();
66        for provider_res in provider_iter {
67            let (id, mut provider) = provider_res.map_err(|e| AppError::Database(e.to_string()))?;
68            provider.id = id.clone();
69
70            // 加载 endpoints
71            let mut stmt_endpoints = conn.prepare(
72                "SELECT url, added_at FROM provider_endpoints WHERE provider_id = ?1 AND app_type = ?2 ORDER BY added_at ASC, url ASC"
73            ).map_err(|e| AppError::Database(e.to_string()))?;
74
75            let endpoints_iter = stmt_endpoints
76                .query_map(params![id, app_type], |row| {
77                    let url: String = row.get(0)?;
78                    let added_at: Option<i64> = row.get(1)?;
79                    Ok((
80                        url,
81                        crate::settings::CustomEndpoint {
82                            url: "".to_string(),
83                            added_at: added_at.unwrap_or(0),
84                            last_used: None,
85                        },
86                    ))
87                })
88                .map_err(|e| AppError::Database(e.to_string()))?;
89
90            let mut custom_endpoints = HashMap::new();
91            for ep_res in endpoints_iter {
92                let (url, mut ep) = ep_res.map_err(|e| AppError::Database(e.to_string()))?;
93                ep.url = url.clone();
94                custom_endpoints.insert(url, ep);
95            }
96
97            if let Some(meta) = &mut provider.meta {
98                meta.custom_endpoints = custom_endpoints;
99            }
100
101            providers.insert(id, provider);
102        }
103
104        Ok(providers)
105    }
106
107    /// 获取当前激活的供应商 ID
108    pub fn get_current_provider(&self, app_type: &str) -> Result<Option<String>, AppError> {
109        let conn = lock_conn!(self.conn);
110        let mut stmt = conn
111            .prepare("SELECT id FROM providers WHERE app_type = ?1 AND is_current = 1 LIMIT 1")
112            .map_err(|e| AppError::Database(e.to_string()))?;
113
114        let mut rows = stmt
115            .query(params![app_type])
116            .map_err(|e| AppError::Database(e.to_string()))?;
117
118        if let Some(row) = rows.next().map_err(|e| AppError::Database(e.to_string()))? {
119            Ok(Some(
120                row.get(0).map_err(|e| AppError::Database(e.to_string()))?,
121            ))
122        } else {
123            Ok(None)
124        }
125    }
126
127    /// 根据 ID 获取单个供应商
128    pub fn get_provider_by_id(
129        &self,
130        id: &str,
131        app_type: &str,
132    ) -> Result<Option<Provider>, AppError> {
133        let conn = lock_conn!(self.conn);
134        let result = conn.query_row(
135            "SELECT name, settings_config, website_url, category, created_at, sort_index, notes, icon, icon_color, meta, in_failover_queue
136             FROM providers WHERE id = ?1 AND app_type = ?2",
137            params![id, app_type],
138            |row| {
139                let name: String = row.get(0)?;
140                let settings_config_str: String = row.get(1)?;
141                let website_url: Option<String> = row.get(2)?;
142                let category: Option<String> = row.get(3)?;
143                let created_at: Option<i64> = row.get(4)?;
144                let sort_index: Option<usize> = row.get(5)?;
145                let notes: Option<String> = row.get(6)?;
146                let icon: Option<String> = row.get(7)?;
147                let icon_color: Option<String> = row.get(8)?;
148                let meta_str: String = row.get(9)?;
149                let in_failover_queue: bool = row.get(10)?;
150
151                let settings_config = serde_json::from_str(&settings_config_str).unwrap_or(serde_json::Value::Null);
152                let meta: ProviderMeta = serde_json::from_str(&meta_str).unwrap_or_default();
153
154                Ok(Provider {
155                    id: id.to_string(),
156                    name,
157                    settings_config,
158                    website_url,
159                    category,
160                    created_at,
161                    sort_index,
162                    notes,
163                    meta: Some(meta),
164                    icon,
165                    icon_color,
166                    in_failover_queue,
167                })
168            },
169        );
170
171        match result {
172            Ok(provider) => Ok(Some(provider)),
173            Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
174            Err(e) => Err(AppError::Database(e.to_string())),
175        }
176    }
177
178    /// 仅获取指定 app 下所有 provider 的 id 集合。
179    pub fn get_provider_ids(&self, app_type: &str) -> Result<HashSet<String>, AppError> {
180        let conn = lock_conn!(self.conn);
181        let mut stmt = conn
182            .prepare("SELECT id FROM providers WHERE app_type = ?1")
183            .map_err(|e| AppError::Database(e.to_string()))?;
184        let rows = stmt
185            .query_map(params![app_type], |row| row.get::<_, String>(0))
186            .map_err(|e| AppError::Database(e.to_string()))?;
187        let mut ids = HashSet::new();
188        for row in rows {
189            ids.insert(row.map_err(|e| AppError::Database(e.to_string()))?);
190        }
191        Ok(ids)
192    }
193
194    /// 判断指定 app 下是否存在非官方种子的供应商。
195    pub fn has_non_official_seed_provider(&self, app_type: &str) -> Result<bool, AppError> {
196        let conn = lock_conn!(self.conn);
197        let mut stmt = conn
198            .prepare("SELECT id FROM providers WHERE app_type = ?1")
199            .map_err(|e| AppError::Database(e.to_string()))?;
200        let mut rows = stmt
201            .query(params![app_type])
202            .map_err(|e| AppError::Database(e.to_string()))?;
203        while let Some(row) = rows.next().map_err(|e| AppError::Database(e.to_string()))? {
204            let id: String = row.get(0).map_err(|e| AppError::Database(e.to_string()))?;
205            if !is_official_seed_id(&id) {
206                return Ok(true);
207            }
208        }
209        Ok(false)
210    }
211
212    fn next_sort_index_for_app(&self, app_type: &str) -> Result<usize, AppError> {
213        let conn = lock_conn!(self.conn);
214        let max: Option<i64> = conn
215            .query_row(
216                "SELECT MAX(sort_index) FROM providers WHERE app_type = ?1",
217                params![app_type],
218                |row| row.get(0),
219            )
220            .map_err(|e| AppError::Database(e.to_string()))?;
221        Ok(max.map(|value| (value + 1) as usize).unwrap_or(0))
222    }
223
224    /// 启动时补齐上游官方预设供应商(Claude / Codex / Gemini)。
225    pub fn init_default_official_providers(&self) -> Result<usize, AppError> {
226        if self
227            .get_bool_flag("official_providers_seeded")
228            .unwrap_or(false)
229        {
230            return Ok(0);
231        }
232
233        let mut inserted = 0usize;
234        let now_ms = chrono::Utc::now().timestamp_millis();
235
236        for seed in OFFICIAL_SEEDS {
237            let app_type = seed.app_type.as_str();
238            if self.get_provider_by_id(seed.id, app_type)?.is_some() {
239                continue;
240            }
241
242            let settings_config: serde_json::Value =
243                serde_json::from_str(seed.settings_config_json).map_err(|err| {
244                    AppError::Database(format!("Seed JSON parse failed for {}: {err}", seed.id))
245                })?;
246
247            let mut provider = Provider::with_id(
248                seed.id.to_string(),
249                seed.name.to_string(),
250                settings_config,
251                Some(seed.website_url.to_string()),
252            );
253            provider.category = Some("official".to_string());
254            provider.icon = Some(seed.icon.to_string());
255            provider.icon_color = Some(seed.icon_color.to_string());
256            provider.sort_index = Some(self.next_sort_index_for_app(app_type)?);
257            provider.created_at = Some(now_ms);
258
259            self.save_provider(app_type, &provider)?;
260            inserted += 1;
261            log::info!("✓ Seeded official provider: {} ({})", seed.name, app_type);
262        }
263
264        self.set_setting("official_providers_seeded", "true")?;
265        Ok(inserted)
266    }
267
268    /// 保存供应商(新增或更新)
269    ///
270    /// 注意:更新模式下不同步 endpoints,因为编辑模式下端点通过单独的 API 管理
271    /// (add_custom_endpoint / remove_custom_endpoint),避免覆盖用户的修改。
272    pub fn save_provider(&self, app_type: &str, provider: &Provider) -> Result<(), AppError> {
273        let mut conn = lock_conn!(self.conn);
274        let tx = conn
275            .transaction()
276            .map_err(|e| AppError::Database(e.to_string()))?;
277
278        // 处理 meta:取出 endpoints 以便单独处理
279        let mut meta_clone = provider.meta.clone().unwrap_or_default();
280        let endpoints = std::mem::take(&mut meta_clone.custom_endpoints);
281
282        // 检查是否存在(用于判断新增/更新,以及保留 is_current 和 in_failover_queue)
283        let existing: Option<(bool, bool)> = tx
284            .query_row(
285                "SELECT is_current, in_failover_queue FROM providers WHERE id = ?1 AND app_type = ?2",
286                params![provider.id, app_type],
287                |row| Ok((row.get(0)?, row.get(1)?)),
288            )
289            .ok();
290
291        let is_update = existing.is_some();
292        let (is_current, in_failover_queue) =
293            existing.unwrap_or((false, provider.in_failover_queue));
294
295        if is_update {
296            // 更新模式:使用 UPDATE 避免触发 ON DELETE CASCADE
297            tx.execute(
298                "UPDATE providers SET
299                    name = ?1,
300                    settings_config = ?2,
301                    website_url = ?3,
302                    category = ?4,
303                    created_at = ?5,
304                    sort_index = ?6,
305                    notes = ?7,
306                    icon = ?8,
307                    icon_color = ?9,
308                    meta = ?10,
309                    is_current = ?11,
310                    in_failover_queue = ?12
311                WHERE id = ?13 AND app_type = ?14",
312                params![
313                    provider.name,
314                    serde_json::to_string(&provider.settings_config).map_err(|e| {
315                        AppError::Database(format!("Failed to serialize settings_config: {e}"))
316                    })?,
317                    provider.website_url,
318                    provider.category,
319                    provider.created_at,
320                    provider.sort_index,
321                    provider.notes,
322                    provider.icon,
323                    provider.icon_color,
324                    serde_json::to_string(&meta_clone).map_err(|e| AppError::Database(format!(
325                        "Failed to serialize meta: {e}"
326                    )))?,
327                    is_current,
328                    in_failover_queue,
329                    provider.id,
330                    app_type,
331                ],
332            )
333            .map_err(|e| AppError::Database(e.to_string()))?;
334        } else {
335            // 新增模式:使用 INSERT
336            tx.execute(
337                "INSERT INTO providers (
338                    id, app_type, name, settings_config, website_url, category,
339                    created_at, sort_index, notes, icon, icon_color, meta, is_current, in_failover_queue
340                ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)",
341                params![
342                    provider.id,
343                    app_type,
344                    provider.name,
345                    serde_json::to_string(&provider.settings_config)
346                        .map_err(|e| AppError::Database(format!("Failed to serialize settings_config: {e}")))?,
347                    provider.website_url,
348                    provider.category,
349                    provider.created_at,
350                    provider.sort_index,
351                    provider.notes,
352                    provider.icon,
353                    provider.icon_color,
354                    serde_json::to_string(&meta_clone)
355                        .map_err(|e| AppError::Database(format!("Failed to serialize meta: {e}")))?,
356                    is_current,
357                    in_failover_queue,
358                ],
359            )
360            .map_err(|e| AppError::Database(e.to_string()))?;
361
362            // 只有新增时才同步 endpoints
363            for (url, endpoint) in endpoints {
364                tx.execute(
365                    "INSERT INTO provider_endpoints (provider_id, app_type, url, added_at)
366                     VALUES (?1, ?2, ?3, ?4)",
367                    params![provider.id, app_type, url, endpoint.added_at],
368                )
369                .map_err(|e| AppError::Database(e.to_string()))?;
370            }
371        }
372
373        tx.commit().map_err(|e| AppError::Database(e.to_string()))?;
374        Ok(())
375    }
376
377    /// 删除供应商
378    pub fn delete_provider(&self, app_type: &str, id: &str) -> Result<(), AppError> {
379        let conn = lock_conn!(self.conn);
380        conn.execute(
381            "DELETE FROM providers WHERE id = ?1 AND app_type = ?2",
382            params![id, app_type],
383        )
384        .map_err(|e| AppError::Database(e.to_string()))?;
385        Ok(())
386    }
387
388    /// 设置当前供应商
389    pub fn set_current_provider(&self, app_type: &str, id: &str) -> Result<(), AppError> {
390        let mut conn = lock_conn!(self.conn);
391        let tx = conn
392            .transaction()
393            .map_err(|e| AppError::Database(e.to_string()))?;
394
395        // 重置所有为 0
396        tx.execute(
397            "UPDATE providers SET is_current = 0 WHERE app_type = ?1",
398            params![app_type],
399        )
400        .map_err(|e| AppError::Database(e.to_string()))?;
401
402        // 设置新的当前供应商
403        tx.execute(
404            "UPDATE providers SET is_current = 1 WHERE id = ?1 AND app_type = ?2",
405            params![id, app_type],
406        )
407        .map_err(|e| AppError::Database(e.to_string()))?;
408
409        tx.commit().map_err(|e| AppError::Database(e.to_string()))?;
410        Ok(())
411    }
412
413    /// 更新供应商的 settings_config(仅更新配置,不改变其他字段)
414    pub fn update_provider_settings_config(
415        &self,
416        app_type: &str,
417        provider_id: &str,
418        settings_config: &serde_json::Value,
419    ) -> Result<(), AppError> {
420        let conn = lock_conn!(self.conn);
421        conn.execute(
422            "UPDATE providers SET settings_config = ?1 WHERE id = ?2 AND app_type = ?3",
423            params![
424                serde_json::to_string(settings_config).map_err(|e| AppError::Database(format!(
425                    "Failed to serialize settings_config: {e}"
426                )))?,
427                provider_id,
428                app_type
429            ],
430        )
431        .map_err(|e| AppError::Database(e.to_string()))?;
432        Ok(())
433    }
434
435    /// 添加自定义端点
436    pub fn add_custom_endpoint(
437        &self,
438        app_type: &str,
439        provider_id: &str,
440        url: &str,
441    ) -> Result<(), AppError> {
442        let conn = lock_conn!(self.conn);
443        let added_at = chrono::Utc::now().timestamp_millis();
444        conn.execute(
445            "INSERT INTO provider_endpoints (provider_id, app_type, url, added_at) VALUES (?1, ?2, ?3, ?4)",
446            params![provider_id, app_type, url, added_at],
447        ).map_err(|e| AppError::Database(e.to_string()))?;
448        Ok(())
449    }
450
451    /// 移除自定义端点
452    pub fn remove_custom_endpoint(
453        &self,
454        app_type: &str,
455        provider_id: &str,
456        url: &str,
457    ) -> Result<(), AppError> {
458        let conn = lock_conn!(self.conn);
459        conn.execute(
460            "DELETE FROM provider_endpoints WHERE provider_id = ?1 AND app_type = ?2 AND url = ?3",
461            params![provider_id, app_type, url],
462        )
463        .map_err(|e| AppError::Database(e.to_string()))?;
464        Ok(())
465    }
466}