Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: make worker identity migrations dry-runnable

Browse Source
This commit is contained in:
Hare
2026-08-20 07:21:03 +09:00
parent de72afd9a1
commit d052cedc7d
8 changed files with 705 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/workspace-server/src/main.rs
+60 -3
View File
@@ -23,6 +23,7 @@ enum Command {
ConfigDiff(WorkspacePathOptions),
Identity(Vec<String>),
TrustRuntime(Vec<String>),
MigrateDryRun { database: Option<PathBuf> },
Skills(SkillsCommand),
Help,
}
@@ -85,6 +86,17 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
Command::ConfigDiff(options) => run_config_diff(options),
Command::Identity(args) => run_identity_command(args),
Command::TrustRuntime(args) => run_trust_runtime_command(args),
Command::MigrateDryRun { database } => {
let database = database.unwrap_or_else(ServerConfig::default_server_database_path);
let plan = SqliteWorkspaceStore::migration_plan(&database).map_err(|error| {
CliError(format!(
"migration dry-run failed for {}: {error}",
database.display()
))
})?;
println!("{}", serde_json::to_string_pretty(&plan)?);
Ok(())
}
Command::Skills(command) => run_skills(command),
Command::Help => Ok(()),
}
@@ -107,6 +119,7 @@ fn parse_command(args: &[String]) -> Result<Command, CliError> {
"config" => parse_config_command(rest),
"identity" => Ok(Command::Identity(rest.to_vec())),
"trust-runtime" => Ok(Command::TrustRuntime(rest.to_vec())),
"migrate" => parse_migrate_command(rest),
"skills" => parse_skills_command(rest),
"serve" => {
if rest.iter().any(|arg| arg == "--help" || arg == "-h") {
@@ -120,7 +133,7 @@ fn parse_command(args: &[String]) -> Result<Command, CliError> {
Ok(Command::Help)
}
other => Err(CliError(format!(
"unknown command `{other}`; expected `init`, `config`, `identity`, `trust-runtime`, `skills`, or `serve`"
"unknown command `{other}`; expected `init`, `config`, `identity`, `trust-runtime`, `migrate`, `skills`, or `serve`"
))),
}
}
@@ -718,6 +731,32 @@ fn parse_config_command(args: &[String]) -> Result<Command, CliError> {
}
}
fn parse_migrate_command(args: &[String]) -> Result<Command, CliError> {
let mut dry_run = false;
let mut database = None;
let mut index = 0;
while index < args.len() {
match args[index].as_str() {
"--dry-run" => dry_run = true,
"--database" => {
index += 1;
database =
Some(PathBuf::from(args.get(index).ok_or_else(|| {
CliError("--database requires a path".to_string())
})?));
}
value => {
return Err(CliError(format!("unknown migrate option: {value}")));
}
}
index += 1;
}
if !dry_run {
return Err(CliError("migrate currently requires --dry-run".to_string()));
}
Ok(Command::MigrateDryRun { database })
}
fn parse_skills_command(args: &[String]) -> Result<Command, CliError> {
let Some((subcommand, rest)) = args.split_first() else {
print_skills_help();
@@ -875,7 +914,8 @@ fn parse_listen(value: &str) -> Result<SocketAddr, CliError> {
fn print_help() {
println!(
"yoi-server\n\nUsage:\n yoi-server init [OPTIONS]\n yoi-server config <COMMAND> [OPTIONS]\n yoi-server identity init --server-id <SERVER_ID> [--replace]\n yoi-server identity show [--json]\n yoi-server trust-runtime add --runtime-id <RUNTIME_ID> --base-url <URL> --public-key <KEY> [--display-name <NAME>] [--replace]\n yoi-server trust-runtime list [--json] [--include-revoked]\n yoi-server trust-runtime revoke --runtime-id <RUNTIME_ID>\n yoi-server skills <COMMAND> [OPTIONS]\n yoi-server serve [OPTIONS]\n\nOptions:\n -h, --help Print help"
"yoi-server\n\nUsage:\n yoi-server init [OPTIONS]\n yoi-server config <COMMAND> [OPTIONS]\n yoi-server identity init --server-id <SERVER_ID> [--replace]\n yoi-server identity show [--json]\n yoi-server trust-runtime add --runtime-id <RUNTIME_ID> --base-url <URL> --public-key <KEY> [--display-name <NAME>] [--replace]\n yoi-server trust-runtime list [--json] [--include-revoked]\n yoi-server trust-runtime revoke --runtime-id <RUNTIME_ID>\n yoi-server skills <COMMAND> [OPTIONS]\n yoi-server migrate --dry-run [--database <PATH>]
yoi-server serve [OPTIONS]\n\nOptions:\n -h, --help Print help"
);
}
@@ -899,7 +939,8 @@ fn print_skills_help() {
fn print_serve_help() {
println!(
"yoi-server serve\n\nUsage:\n yoi-server serve [OPTIONS]\n\nDescription:\n Serves the Workspace recorded in the Yoi server DB. Workspace records are stored in the XDG/Yoi data directory, and runtime sources are loaded from XDG runtimes.toml.\n\nOptions:\n --listen <ADDR> Listen address (default 127.0.0.1:8787)\n -h, --help Print help"
"yoi-server serve\n\nUsage:\n yoi-server migrate --dry-run [--database <PATH>]
yoi-server serve [OPTIONS]\n\nDescription:\n Serves the Workspace recorded in the Yoi server DB. Workspace records are stored in the XDG/Yoi data directory, and runtime sources are loaded from XDG runtimes.toml.\n\nOptions:\n --listen <ADDR> Listen address (default 127.0.0.1:8787)\n -h, --help Print help"
);
}
@@ -939,6 +980,22 @@ mod tests {
assert_eq!(name, "debug-rust");
}
#[test]
fn parse_migrate_requires_dry_run_and_accepts_database_path() {
let error = parse_migrate_command(&[]).unwrap_err();
assert_eq!(error.to_string(), "migrate currently requires --dry-run");
let command = parse_migrate_command(&[
"--dry-run".to_string(),
"--database".to_string(),
"/tmp/server.db".to_string(),
])
.unwrap();
let Command::MigrateDryRun { database } = command else {
panic!("expected migration dry-run command");
};
assert_eq!(database, Some(PathBuf::from("/tmp/server.db")));
}
#[test]
fn parse_serve_accepts_listen_only() {
let args = vec!["--listen".to_string(), "127.0.0.1:0".to_string()];
crates/workspace-server/src/retention.rs
+19
View File
@@ -166,6 +166,25 @@ pub enum WorkerRetentionError {
Invalid(String),
}
pub(crate) fn repair_worker_diagnostics_archive_table(conn: &Connection) -> crate::Result<bool> {
let existed: bool = conn.query_row(
"SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type='table' AND name='worker_diagnostics_archives')",
[],
|row| row.get(0),
)?;
if !existed {
conn.execute_batch(
"CREATE TABLE worker_diagnostics_archives (
operation_id TEXT PRIMARY KEY, workspace_id TEXT NOT NULL, runtime_id TEXT NOT NULL,
worker_id TEXT NOT NULL, policy_id TEXT NOT NULL, policy_revision INTEGER NOT NULL,
committed_at TEXT NOT NULL, expires_at TEXT NOT NULL,
FOREIGN KEY(operation_id) REFERENCES worker_removal_operations(operation_id),
FOREIGN KEY(workspace_id) REFERENCES workspaces(workspace_id) ON DELETE CASCADE);",
)?;
}
Ok(!existed)
}
pub(crate) fn create_worker_retention_tables(conn: &Connection) -> crate::Result<()> {
conn.execute_batch(r#"
CREATE TABLE workspace_worker_retention_policy_revisions (
crates/workspace-server/src/store.rs
+141 -8
View File
@@ -4,11 +4,15 @@ use std::time::Duration;
use async_trait::async_trait;
use flow::{CompiledFlowDefinition, FlowSourceKind, compile_flow_source};
use rusqlite::{Connection, OptionalExtension, TransactionBehavior, params};
use rusqlite::{
Connection, OpenFlags, OptionalExtension, TransactionBehavior, backup::Backup, params,
};
use serde::{Deserialize, Serialize};
use uuid::Uuid;
use worker_runtime::identity::{RuntimeWorkerRef, WorkerId};
use worker_runtime::identity::{
LegacyWorkerIdentityMapping, RuntimeWorkerRef, WorkerId, legacy_worker_identity_mapping_digest,
};
use crate::{Error, Result};
@@ -204,7 +208,7 @@ const MIGRATIONS: &[Migration] = &[
Migration {
version: 37,
name: "promote Workspace Worker UUIDv7 identity",
apply: promote_workspace_worker_uuid_identity,
apply: apply_workspace_worker_uuid_identity_migration,
},
Migration {
version: 38,
@@ -219,6 +223,17 @@ struct Migration {
apply: fn(&Connection) -> Result<()>,
}
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct WorkspaceStoreMigrationPlan {
pub current_schema_version: i64,
pub target_schema_version: i64,
pub migration_required: bool,
pub worker_count: usize,
pub mapping_digest: String,
pub mappings: Vec<LegacyWorkerIdentityMapping>,
pub repairs: Vec<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct WorkspaceRecord {
pub workspace_id: String,
@@ -976,6 +991,66 @@ pub struct SqliteWorkspaceStore {
}
impl SqliteWorkspaceStore {
pub fn migration_plan(path: impl AsRef<Path>) -> Result<WorkspaceStoreMigrationPlan> {
let path = path.as_ref();
let source = Connection::open_with_flags(
path,
OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_URI,
)?;
let current_schema_version = current_schema_version(&source)?;
let target_schema_version = MIGRATIONS
.last()
.map(|migration| i64::from(migration.version))
.unwrap_or(current_schema_version);
let mut repairs = Vec::new();
if current_schema_version < 37 && !table_exists(&source, "worker_diagnostics_archives")? {
repairs.push("create missing worker_diagnostics_archives table".to_string());
}
let mut candidate = Connection::open_in_memory()?;
{
let backup = Backup::new(&source, &mut candidate)?;
backup.run_to_completion(5, Duration::from_millis(10), None)?;
}
configure_sqlite(&candidate)?;
let mappings = if current_schema_version < 37 {
apply_migrations_through(&candidate, 36)?;
let tx = candidate.unchecked_transaction()?;
crate::retention::repair_worker_diagnostics_archive_table(&tx)?;
let mappings = promote_workspace_worker_uuid_identity(&tx)?;
tx.execute(
"INSERT INTO __yoi_schema_migrations (version, name) VALUES (37, ?1)",
["promote Workspace Worker UUIDv7 identity"],
)?;
tx.commit()?;
mappings
} else {
Vec::new()
};
apply_migrations_through(&candidate, i64::MAX)?;
ticket::migrate_sqlite_ticket_schema(&candidate)?;
merge_request::migrate(&candidate).map_err(|error| Error::Store(error.to_string()))?;
validate_workspace_repository_references(&candidate)?;
let foreign_key_failures: i64 =
candidate.query_row("SELECT COUNT(*) FROM pragma_foreign_key_check", [], |row| {
row.get(0)
})?;
if foreign_key_failures != 0 {
return Err(Error::Store(format!(
"migration dry-run found {foreign_key_failures} foreign key violation(s)"
)));
}
Ok(WorkspaceStoreMigrationPlan {
current_schema_version,
target_schema_version,
migration_required: current_schema_version < target_schema_version,
worker_count: mappings.len(),
mapping_digest: legacy_worker_identity_mapping_digest(&mappings),
mappings,
repairs,
})
}
pub fn open(path: impl AsRef<Path>) -> Result<Self> {
let conn = Connection::open(path)?;
Self::from_connection(conn)
@@ -5310,7 +5385,13 @@ fn collect_legacy_text_worker_bindings(
Ok(())
}
fn promote_workspace_worker_uuid_identity(conn: &Connection) -> Result<()> {
fn apply_workspace_worker_uuid_identity_migration(conn: &Connection) -> Result<()> {
promote_workspace_worker_uuid_identity(conn).map(|_| ())
}
fn promote_workspace_worker_uuid_identity(
conn: &Connection,
) -> Result<Vec<LegacyWorkerIdentityMapping>> {
conn.execute_batch(
r#"
PRAGMA defer_foreign_keys = ON;
@@ -5367,7 +5448,9 @@ fn promote_workspace_worker_uuid_identity(conn: &Connection) -> Result<()> {
collect_legacy_text_worker_bindings(conn, table, &mut legacy_workers)?;
}
for (workspace_id, runtime_id, runtime_worker_id) in legacy_workers {
let mut mappings = Vec::with_capacity(legacy_workers.len());
for (workspace_id, runtime_id, runtime_worker_id) in &legacy_workers {
let worker_id = WorkerId::from_legacy_binding(workspace_id, runtime_id, *runtime_worker_id);
conn.execute(
"INSERT INTO worker_identity_v37(\
workspace_id, runtime_id, runtime_worker_id, worker_id\
@@ -5376,10 +5459,15 @@ fn promote_workspace_worker_uuid_identity(conn: &Connection) -> Result<()> {
workspace_id,
runtime_id,
runtime_worker_id,
WorkerId::from_legacy_binding(&workspace_id, &runtime_id, runtime_worker_id)
.to_string()
worker_id.to_string()
],
)?;
mappings.push(LegacyWorkerIdentityMapping {
workspace_id: workspace_id.clone(),
runtime_id: runtime_id.clone(),
legacy_worker_id: *runtime_worker_id,
worker_id,
});
}
conn.execute_batch(
@@ -5581,7 +5669,7 @@ fn promote_workspace_worker_uuid_identity(conn: &Connection) -> Result<()> {
DROP TABLE worker_identity_v37;
"#,
)?;
Ok(())
Ok(mappings)
}
fn allocate_resource_human_key(
@@ -5852,6 +5940,9 @@ pub(crate) fn apply_migrations_through(conn: &Connection, through_version: i64)
i64::from(migration.version) > current && i64::from(migration.version) <= through_version
}) {
let tx = conn.unchecked_transaction()?;
if migration.version == 37 {
crate::retention::repair_worker_diagnostics_archive_table(&tx)?;
}
(migration.apply)(&tx)?;
tx.execute(
"INSERT INTO __yoi_schema_migrations (version, name) VALUES (?1, ?2)",
@@ -6369,6 +6460,48 @@ mod tests {
.unwrap();
}
#[test]
fn migration_dry_run_repairs_missing_diagnostics_archive_without_mutating_source() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().join("server.db");
{
let conn = Connection::open(&path).unwrap();
configure_sqlite(&conn).unwrap();
apply_migrations_through(&conn, 36).unwrap();
conn.execute_batch(
"DROP TABLE worker_diagnostics_archives;
INSERT INTO workspaces(workspace_id, display_name, state, created_at, updated_at)
VALUES ('workspace-a', 'Workspace A', 'active', '1', '1');
INSERT INTO worker_registry(
workspace_id, runtime_id, runtime_worker_id, display_name,
retention_state, created_at, updated_at
) VALUES ('workspace-a', 'runtime-a', 7, 'Worker 7', 'normal', '1', '1');",
)
.unwrap();
}
let before = std::fs::read(&path).unwrap();
let plan = SqliteWorkspaceStore::migration_plan(&path).unwrap();
assert_eq!(plan.current_schema_version, 36);
assert_eq!(plan.target_schema_version, 38);
assert!(plan.migration_required);
assert_eq!(plan.worker_count, 1);
assert_eq!(plan.mappings[0].legacy_worker_id, 7);
assert_eq!(
plan.repairs,
vec!["create missing worker_diagnostics_archives table"]
);
assert_eq!(std::fs::read(&path).unwrap(), before);
let store = SqliteWorkspaceStore::open(&path).unwrap();
store
.with_conn(|conn| {
assert!(table_exists(conn, "worker_diagnostics_archives")?);
assert_eq!(current_schema_version(conn)?, 38);
Ok(())
})
.unwrap();
}
#[test]
fn v38_backfills_workspace_scoped_objective_and_worker_human_keys() {
let conn = Connection::open_in_memory().unwrap();