forked from boos/BoosAPI
feat: add workflow studio and postgres support
This commit is contained in:
+100
-1
@@ -1,12 +1,111 @@
|
||||
use sqlx::sqlite::{SqlitePool, SqlitePoolOptions};
|
||||
use anyhow::{bail, Context};
|
||||
use serde::Serialize;
|
||||
use sqlx::{
|
||||
any::{install_default_drivers, AnyPoolOptions},
|
||||
postgres::{PgPool, PgPoolOptions},
|
||||
sqlite::{SqlitePool, SqlitePoolOptions},
|
||||
};
|
||||
|
||||
pub mod models;
|
||||
pub mod queries;
|
||||
pub mod sqlite_to_postgres;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum DbKind {
|
||||
Sqlite,
|
||||
Postgres,
|
||||
}
|
||||
|
||||
impl DbKind {
|
||||
pub fn from_url(database_url: &str) -> anyhow::Result<Self> {
|
||||
let normalized = database_url.trim().to_ascii_lowercase();
|
||||
if normalized.starts_with("sqlite:") {
|
||||
Ok(Self::Sqlite)
|
||||
} else if normalized.starts_with("postgres://") || normalized.starts_with("postgresql://") {
|
||||
Ok(Self::Postgres)
|
||||
} else {
|
||||
bail!("unsupported database URL; expected sqlite:, postgres:// or postgresql://")
|
||||
}
|
||||
}
|
||||
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Sqlite => "sqlite",
|
||||
Self::Postgres => "postgres",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn create_pool(database_url: &str, max_connections: u32) -> anyhow::Result<SqlitePool> {
|
||||
if DbKind::from_url(database_url)? != DbKind::Sqlite {
|
||||
bail!(
|
||||
"runtime PostgreSQL query support is not enabled yet; use migrate-sqlite-to-postgres to prepare PostgreSQL data"
|
||||
);
|
||||
}
|
||||
let pool = SqlitePoolOptions::new()
|
||||
.max_connections(max_connections)
|
||||
.connect(database_url)
|
||||
.await?;
|
||||
Ok(pool)
|
||||
}
|
||||
|
||||
pub async fn create_pg_pool(database_url: &str, max_connections: u32) -> anyhow::Result<PgPool> {
|
||||
let pool = PgPoolOptions::new()
|
||||
.max_connections(max_connections)
|
||||
.connect(database_url)
|
||||
.await
|
||||
.context("connect PostgreSQL")?;
|
||||
Ok(pool)
|
||||
}
|
||||
|
||||
pub async fn test_connection(database_url: &str, max_connections: u32) -> anyhow::Result<DbKind> {
|
||||
match DbKind::from_url(database_url)? {
|
||||
DbKind::Sqlite => {
|
||||
SqlitePoolOptions::new()
|
||||
.max_connections(max_connections)
|
||||
.connect(database_url)
|
||||
.await
|
||||
.context("connect SQLite")?
|
||||
.close()
|
||||
.await;
|
||||
Ok(DbKind::Sqlite)
|
||||
}
|
||||
DbKind::Postgres => {
|
||||
create_pg_pool(database_url, max_connections)
|
||||
.await?
|
||||
.close()
|
||||
.await;
|
||||
Ok(DbKind::Postgres)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run_migrations(
|
||||
kind: DbKind,
|
||||
sqlite_pool: Option<&SqlitePool>,
|
||||
pg_pool: Option<&PgPool>,
|
||||
) -> anyhow::Result<()> {
|
||||
match kind {
|
||||
DbKind::Sqlite => {
|
||||
let pool = sqlite_pool.context("SQLite pool required for SQLite migrations")?;
|
||||
sqlx::migrate!("./migrations").run(pool).await?;
|
||||
}
|
||||
DbKind::Postgres => {
|
||||
let pool = pg_pool.context("PostgreSQL pool required for PostgreSQL migrations")?;
|
||||
sqlx::migrate!("./migrations_pg").run(pool).await?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn create_any_pool(
|
||||
database_url: &str,
|
||||
max_connections: u32,
|
||||
) -> anyhow::Result<sqlx::AnyPool> {
|
||||
install_default_drivers();
|
||||
Ok(AnyPoolOptions::new()
|
||||
.max_connections(max_connections)
|
||||
.connect(database_url)
|
||||
.await?)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,463 @@
|
||||
use anyhow::{bail, Context};
|
||||
use serde::Serialize;
|
||||
use sqlx::{
|
||||
postgres::PgPool,
|
||||
sqlite::{SqlitePool, SqliteRow},
|
||||
Row,
|
||||
};
|
||||
use std::collections::HashMap;
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct MigrationSummary {
|
||||
pub tables: Vec<TableMigrationSummary>,
|
||||
pub total_rows: i64,
|
||||
pub force_truncate: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct TableMigrationSummary {
|
||||
pub table: &'static str,
|
||||
pub rows: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
enum ColumnKind {
|
||||
I64,
|
||||
F64,
|
||||
Bool,
|
||||
Text,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct ColumnSpec {
|
||||
name: &'static str,
|
||||
kind: ColumnKind,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct TableSpec {
|
||||
name: &'static str,
|
||||
columns: &'static [ColumnSpec],
|
||||
}
|
||||
|
||||
macro_rules! col {
|
||||
($name:literal, i64) => {
|
||||
ColumnSpec {
|
||||
name: $name,
|
||||
kind: ColumnKind::I64,
|
||||
}
|
||||
};
|
||||
($name:literal, f64) => {
|
||||
ColumnSpec {
|
||||
name: $name,
|
||||
kind: ColumnKind::F64,
|
||||
}
|
||||
};
|
||||
($name:literal, bool) => {
|
||||
ColumnSpec {
|
||||
name: $name,
|
||||
kind: ColumnKind::Bool,
|
||||
}
|
||||
};
|
||||
($name:literal, text) => {
|
||||
ColumnSpec {
|
||||
name: $name,
|
||||
kind: ColumnKind::Text,
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
const TABLES: &[TableSpec] = &[
|
||||
TableSpec {
|
||||
name: "users",
|
||||
columns: &[
|
||||
col!("id", i64),
|
||||
col!("username", text),
|
||||
col!("email", text),
|
||||
col!("password_hash", text),
|
||||
col!("role", text),
|
||||
col!("status", text),
|
||||
col!("created_at", text),
|
||||
col!("updated_at", text),
|
||||
],
|
||||
},
|
||||
TableSpec {
|
||||
name: "sites",
|
||||
columns: &[
|
||||
col!("id", i64),
|
||||
col!("name", text),
|
||||
col!("url", text),
|
||||
col!("platform", text),
|
||||
col!("status", text),
|
||||
col!("api_key", text),
|
||||
col!("proxy_url", text),
|
||||
col!("use_system_proxy", bool),
|
||||
col!("custom_headers", text),
|
||||
col!("external_checkin_url", text),
|
||||
col!("global_weight", f64),
|
||||
col!("post_refresh_probe_enabled", bool),
|
||||
col!("post_refresh_probe_model", text),
|
||||
col!("post_refresh_probe_scope", text),
|
||||
col!("post_refresh_probe_latency_threshold_ms", i64),
|
||||
col!("is_pinned", bool),
|
||||
col!("sort_order", i64),
|
||||
col!("created_at", text),
|
||||
col!("updated_at", text),
|
||||
],
|
||||
},
|
||||
TableSpec {
|
||||
name: "accounts",
|
||||
columns: &[
|
||||
col!("id", i64),
|
||||
col!("site_id", i64),
|
||||
col!("username", text),
|
||||
col!("access_token", text),
|
||||
col!("api_token", text),
|
||||
col!("balance", f64),
|
||||
col!("balance_used", f64),
|
||||
col!("quota", f64),
|
||||
col!("unit_cost", f64),
|
||||
col!("value_score", f64),
|
||||
col!("status", text),
|
||||
col!("checkin_enabled", bool),
|
||||
col!("last_checkin_at", text),
|
||||
col!("last_balance_refresh", text),
|
||||
col!("extra_config", text),
|
||||
col!("is_pinned", bool),
|
||||
col!("sort_order", i64),
|
||||
col!("oauth_provider", text),
|
||||
col!("oauth_account_key", text),
|
||||
col!("oauth_project_id", text),
|
||||
col!("created_at", text),
|
||||
col!("updated_at", text),
|
||||
],
|
||||
},
|
||||
TableSpec {
|
||||
name: "account_tokens",
|
||||
columns: &[
|
||||
col!("id", i64),
|
||||
col!("account_id", i64),
|
||||
col!("name", text),
|
||||
col!("token", text),
|
||||
col!("source", text),
|
||||
col!("token_group", text),
|
||||
col!("value_status", text),
|
||||
col!("enabled", bool),
|
||||
col!("is_default", bool),
|
||||
col!("created_at", text),
|
||||
col!("updated_at", text),
|
||||
],
|
||||
},
|
||||
TableSpec {
|
||||
name: "downstream_api_keys",
|
||||
columns: &[
|
||||
col!("id", i64),
|
||||
col!("name", text),
|
||||
col!("key", text),
|
||||
col!("description", text),
|
||||
col!("group_name", text),
|
||||
col!("tags", text),
|
||||
col!("enabled", bool),
|
||||
col!("expires_at", text),
|
||||
col!("max_cost", f64),
|
||||
col!("used_cost", f64),
|
||||
col!("max_requests", i64),
|
||||
col!("used_requests", i64),
|
||||
col!("supported_models", text),
|
||||
col!("allowed_route_ids", text),
|
||||
col!("site_weight_multipliers", text),
|
||||
col!("excluded_site_ids", text),
|
||||
col!("excluded_credential_refs", text),
|
||||
col!("last_used_at", text),
|
||||
col!("user_id", i64),
|
||||
col!("created_at", text),
|
||||
col!("updated_at", text),
|
||||
],
|
||||
},
|
||||
TableSpec {
|
||||
name: "token_routes",
|
||||
columns: &[
|
||||
col!("id", i64),
|
||||
col!("model_pattern", text),
|
||||
col!("model_mapping", text),
|
||||
col!("display_name", text),
|
||||
col!("display_icon", text),
|
||||
col!("decision_snapshot", text),
|
||||
col!("decision_refreshed_at", text),
|
||||
col!("routing_strategy", text),
|
||||
col!("route_mode", text),
|
||||
col!("enabled", bool),
|
||||
col!("created_at", text),
|
||||
col!("updated_at", text),
|
||||
],
|
||||
},
|
||||
TableSpec {
|
||||
name: "route_channels",
|
||||
columns: &[
|
||||
col!("id", i64),
|
||||
col!("route_id", i64),
|
||||
col!("account_id", i64),
|
||||
col!("token_id", i64),
|
||||
col!("oauth_route_unit_id", i64),
|
||||
col!("source_model", text),
|
||||
col!("priority", i64),
|
||||
col!("weight", i64),
|
||||
col!("enabled", bool),
|
||||
col!("manual_override", bool),
|
||||
col!("success_count", i64),
|
||||
col!("fail_count", i64),
|
||||
col!("total_latency_ms", i64),
|
||||
col!("total_cost", f64),
|
||||
col!("last_used_at", text),
|
||||
col!("last_selected_at", text),
|
||||
col!("last_fail_at", text),
|
||||
col!("consecutive_fail_count", i64),
|
||||
col!("cooldown_level", i64),
|
||||
col!("cooldown_until", text),
|
||||
],
|
||||
},
|
||||
TableSpec {
|
||||
name: "proxy_logs",
|
||||
columns: &[
|
||||
col!("id", i64),
|
||||
col!("route_id", i64),
|
||||
col!("channel_id", i64),
|
||||
col!("account_id", i64),
|
||||
col!("downstream_api_key_id", i64),
|
||||
col!("model_requested", text),
|
||||
col!("model_actual", text),
|
||||
col!("status", text),
|
||||
col!("http_status", i64),
|
||||
col!("latency_ms", i64),
|
||||
col!("first_byte_latency_ms", i64),
|
||||
col!("prompt_tokens", i64),
|
||||
col!("completion_tokens", i64),
|
||||
col!("total_tokens", i64),
|
||||
col!("estimated_cost", f64),
|
||||
col!("error_message", text),
|
||||
col!("retry_count", i64),
|
||||
col!("is_stream", bool),
|
||||
col!("client_family", text),
|
||||
col!("client_app_id", text),
|
||||
col!("client_app_name", text),
|
||||
col!("client_confidence", text),
|
||||
col!("billing_details", text),
|
||||
col!("created_at", text),
|
||||
],
|
||||
},
|
||||
TableSpec {
|
||||
name: "settings",
|
||||
columns: &[col!("key", text), col!("value", text)],
|
||||
},
|
||||
TableSpec {
|
||||
name: "workspace_assets",
|
||||
columns: &[
|
||||
col!("id", i64),
|
||||
col!("public_id", text),
|
||||
col!("user_id", i64),
|
||||
col!("source", text),
|
||||
col!("kind", text),
|
||||
col!("title", text),
|
||||
col!("prompt", text),
|
||||
col!("model", text),
|
||||
col!("status", text),
|
||||
col!("mime_type", text),
|
||||
col!("file_name", text),
|
||||
col!("file_path", text),
|
||||
col!("remote_url", text),
|
||||
col!("thumbnail_path", text),
|
||||
col!("size_bytes", i64),
|
||||
col!("metadata", text),
|
||||
col!("created_at", text),
|
||||
col!("updated_at", text),
|
||||
col!("deleted_at", text),
|
||||
],
|
||||
},
|
||||
TableSpec {
|
||||
name: "user_announcements",
|
||||
columns: &[
|
||||
col!("id", i64),
|
||||
col!("title", text),
|
||||
col!("content", text),
|
||||
col!("level", text),
|
||||
col!("status", text),
|
||||
col!("audience", text),
|
||||
col!("starts_at", text),
|
||||
col!("ends_at", text),
|
||||
col!("created_at", text),
|
||||
col!("updated_at", text),
|
||||
col!("published_at", text),
|
||||
],
|
||||
},
|
||||
TableSpec {
|
||||
name: "user_announcement_reads",
|
||||
columns: &[
|
||||
col!("announcement_id", i64),
|
||||
col!("user_id", i64),
|
||||
col!("read_at", text),
|
||||
],
|
||||
},
|
||||
];
|
||||
|
||||
pub async fn migrate(
|
||||
sqlite: &SqlitePool,
|
||||
pg: &PgPool,
|
||||
force_truncate: bool,
|
||||
) -> anyhow::Result<MigrationSummary> {
|
||||
ensure_target_is_empty(pg, force_truncate).await?;
|
||||
if force_truncate {
|
||||
truncate_target(pg).await?;
|
||||
}
|
||||
|
||||
let mut summaries = Vec::new();
|
||||
let mut total_rows = 0;
|
||||
for table in TABLES {
|
||||
let rows = copy_table(sqlite, pg, table).await?;
|
||||
summaries.push(TableMigrationSummary {
|
||||
table: table.name,
|
||||
rows,
|
||||
});
|
||||
total_rows += rows;
|
||||
}
|
||||
sync_sequences(pg).await?;
|
||||
Ok(MigrationSummary {
|
||||
tables: summaries,
|
||||
total_rows,
|
||||
force_truncate,
|
||||
})
|
||||
}
|
||||
|
||||
async fn ensure_target_is_empty(pg: &PgPool, force_truncate: bool) -> anyhow::Result<()> {
|
||||
if force_truncate {
|
||||
return Ok(());
|
||||
}
|
||||
let mut non_empty = Vec::new();
|
||||
for table in TABLES {
|
||||
let sql = format!("SELECT COUNT(*) FROM {}", table.name);
|
||||
let count: i64 = sqlx::query_scalar(&sql).fetch_one(pg).await.unwrap_or(0);
|
||||
if count > 0 {
|
||||
non_empty.push(format!("{} ({})", table.name, count));
|
||||
}
|
||||
}
|
||||
if !non_empty.is_empty() {
|
||||
bail!(
|
||||
"target PostgreSQL database is not empty: {}; rerun with --force-truncate to overwrite",
|
||||
non_empty.join(", ")
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn truncate_target(pg: &PgPool) -> anyhow::Result<()> {
|
||||
let table_names = TABLES
|
||||
.iter()
|
||||
.rev()
|
||||
.map(|table| table.name)
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
let sql = format!("TRUNCATE TABLE {} RESTART IDENTITY CASCADE", table_names);
|
||||
sqlx::query(&sql).execute(pg).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn copy_table(sqlite: &SqlitePool, pg: &PgPool, table: &TableSpec) -> anyhow::Result<i64> {
|
||||
let available = sqlite_columns(sqlite, table.name).await?;
|
||||
let columns = table
|
||||
.columns
|
||||
.iter()
|
||||
.copied()
|
||||
.filter(|column| available.contains_key(column.name))
|
||||
.collect::<Vec<_>>();
|
||||
if columns.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let select_sql = format!(
|
||||
"SELECT {} FROM {} ORDER BY {}",
|
||||
columns
|
||||
.iter()
|
||||
.map(|column| column.name)
|
||||
.collect::<Vec<_>>()
|
||||
.join(", "),
|
||||
table.name,
|
||||
if table.name == "settings" {
|
||||
"key"
|
||||
} else {
|
||||
columns[0].name
|
||||
}
|
||||
);
|
||||
let rows = sqlx::query(&select_sql).fetch_all(sqlite).await?;
|
||||
let insert_sql = format!(
|
||||
"INSERT INTO {} ({}) VALUES ({})",
|
||||
table.name,
|
||||
columns
|
||||
.iter()
|
||||
.map(|column| quote_ident(column.name))
|
||||
.collect::<Vec<_>>()
|
||||
.join(", "),
|
||||
(1..=columns.len())
|
||||
.map(|index| format!("${index}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ")
|
||||
);
|
||||
|
||||
for row in &rows {
|
||||
let mut query = sqlx::query(&insert_sql);
|
||||
for column in &columns {
|
||||
query = bind_column(query, row, column)?;
|
||||
}
|
||||
query
|
||||
.execute(pg)
|
||||
.await
|
||||
.with_context(|| format!("insert into {}", table.name))?;
|
||||
}
|
||||
Ok(rows.len() as i64)
|
||||
}
|
||||
|
||||
fn bind_column<'q>(
|
||||
query: sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments>,
|
||||
row: &SqliteRow,
|
||||
column: &ColumnSpec,
|
||||
) -> anyhow::Result<sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments>> {
|
||||
Ok(match column.kind {
|
||||
ColumnKind::I64 => query.bind(row.try_get::<Option<i64>, _>(column.name)?),
|
||||
ColumnKind::F64 => query.bind(row.try_get::<Option<f64>, _>(column.name)?),
|
||||
ColumnKind::Bool => query.bind(row.try_get::<Option<bool>, _>(column.name)?),
|
||||
ColumnKind::Text => query.bind(row.try_get::<Option<String>, _>(column.name)?),
|
||||
})
|
||||
}
|
||||
|
||||
async fn sqlite_columns(sqlite: &SqlitePool, table: &str) -> anyhow::Result<HashMap<String, ()>> {
|
||||
let rows = sqlx::query(&format!("PRAGMA table_info({table})"))
|
||||
.fetch_all(sqlite)
|
||||
.await?;
|
||||
Ok(rows
|
||||
.into_iter()
|
||||
.filter_map(|row| row.try_get::<String, _>("name").ok())
|
||||
.map(|name| (name, ()))
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn sync_sequences(pg: &PgPool) -> anyhow::Result<()> {
|
||||
for table in TABLES
|
||||
.iter()
|
||||
.filter(|table| table.columns.iter().any(|column| column.name == "id"))
|
||||
{
|
||||
let sql = format!(
|
||||
"SELECT setval(pg_get_serial_sequence('{table}', 'id'), COALESCE((SELECT MAX(id) FROM {table}), 1), (SELECT COUNT(*) FROM {table}) > 0)",
|
||||
table = table.name
|
||||
);
|
||||
let _ = sqlx::query(&sql).execute(pg).await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn quote_ident(name: &str) -> String {
|
||||
if name == "read" {
|
||||
"\"read\"".to_string()
|
||||
} else {
|
||||
name.to_string()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user