feat: add workflow studio and postgres support

This commit is contained in:
boos
2026-05-25 16:03:07 +08:00
parent 7cdb02ef1d
commit cb35d5fc6e
96 changed files with 12414 additions and 15548 deletions
+100 -1
View File
@@ -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?)
}
+463
View File
@@ -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()
}
}