chronicle_proxy/
database.rs

1use std::{collections::BTreeMap, sync::Arc};
2
3use error_stack::Report;
4use logging::ProxyLogEntry;
5
6use crate::{
7    config::{AliasConfig, ApiKeyConfig, CustomProviderConfig},
8    providers::custom::ProviderRequestFormat,
9    Error,
10};
11
12pub mod logging;
13#[cfg(feature = "postgres")]
14pub mod postgres;
15#[cfg(feature = "sqlite")]
16pub mod sqlite;
17#[cfg(test)]
18mod testing;
19
20/// A DBMS-agnostic interface to a database
21#[async_trait::async_trait]
22pub trait ProxyDatabase: std::fmt::Debug + Send + Sync {
23    /// Load provider configuration from the database
24    async fn load_providers_from_database(
25        &self,
26        providers_table: &str,
27    ) -> Result<Vec<DbProvider>, Report<Error>>;
28
29    /// Load alias configuration from the database
30    async fn load_aliases_from_database(
31        &self,
32        alias_table: &str,
33        providers_table: &str,
34    ) -> Result<Vec<AliasConfig>, Report<Error>>;
35
36    /// Load API key configuration from the database
37    async fn load_api_key_configs_from_database(
38        &self,
39        table: &str,
40    ) -> Result<Vec<ApiKeyConfig>, Report<Error>>;
41
42    /// Write a batch of log entries to the database
43    async fn write_log_batch(&self, items: Vec<ProxyLogEntry>) -> Result<(), sqlx::Error>;
44}
45
46/// A [ProxyDatabase] wrapped in an [Arc]
47pub type Database = Arc<dyn ProxyDatabase>;
48
49/// A provider configuration loaded from the database
50#[derive(sqlx::FromRow)]
51pub struct DbProvider {
52    name: String,
53    label: Option<String>,
54    url: String,
55    api_key: Option<String>,
56    format: sqlx::types::Json<ProviderRequestFormat>,
57    headers: Option<sqlx::types::Json<BTreeMap<String, String>>>,
58    prefix: Option<String>,
59    api_key_source: Option<String>,
60}
61
62/// Load provider configuration from the database
63pub async fn load_providers_from_database(
64    db: &dyn ProxyDatabase,
65    providers_table: &str,
66) -> Result<Vec<CustomProviderConfig>, Report<Error>> {
67    let rows = db.load_providers_from_database(providers_table).await?;
68    let providers = rows
69        .into_iter()
70        .map(|row| CustomProviderConfig {
71            name: row.name,
72            label: row.label,
73            url: row.url,
74            api_key: row.api_key,
75            format: row.format.0,
76            headers: row.headers.unwrap_or_default().0,
77            prefix: row.prefix,
78            api_key_source: row.api_key_source,
79        })
80        .collect();
81    Ok(providers)
82}