1use crate::cmd::utils as cmd_utils;
2use crate::pkg::{config, hooks, install, local, pin, resolve, transaction, types};
3use anyhow::{Result, anyhow};
4use colored::*;
5use indicatif::{ProgressBar, ProgressStyle};
6use rayon::prelude::*;
7use semver::Version;
8use std::fs;
9use std::sync::Mutex;
10
11pub fn run(all: bool, package_names: &[String], yes: bool) -> Result<()> {
12 if all {
13 return run_update_all_logic(yes);
14 }
15
16 let expanded_package_names = cmd_utils::expand_split_packages(package_names, "Updating")?;
17
18 let mut failed_packages = Vec::new();
19
20 for (i, package_name) in expanded_package_names.iter().enumerate() {
21 if i > 0 {
22 println!();
23 }
24 if let Err(e) = run_update_single_logic(package_name, yes) {
25 eprintln!(
26 "{}: Failed to update '{}': {}",
27 "Error".red().bold(),
28 package_name,
29 e
30 );
31 failed_packages.push(package_name.clone());
32 }
33 }
34
35 if !failed_packages.is_empty() {
36 return Err(anyhow!(
37 "The following packages failed to update: {}",
38 failed_packages.join(", ")
39 ));
40 } else if !package_names.is_empty() {
41 println!("\n{}", "Success:".green());
42 }
43 Ok(())
44}
45
46fn run_update_single_logic(package_name: &str, yes: bool) -> Result<()> {
47 println!("--- Updating package '{}' ---", package_name.blue().bold());
48
49 let request = resolve::parse_source_string(package_name)?;
50
51 let (new_pkg, new_version, _, _, registry_handle) =
52 resolve::resolve_package_and_version(package_name, true)?;
53
54 if pin::is_pinned(package_name)? {
55 println!(
56 "Package '{}' is pinned. Skipping update.",
57 package_name.yellow()
58 );
59 return Ok(());
60 }
61
62 let old_manifest = match local::is_package_installed(
63 &new_pkg.name,
64 request.sub_package.as_deref(),
65 types::Scope::User,
66 )?
67 .or(local::is_package_installed(
68 &new_pkg.name,
69 request.sub_package.as_deref(),
70 types::Scope::System,
71 )?) {
72 Some(m) => m,
73 None => {
74 return Err(anyhow!(
75 "Package '{package_name}' is not installed. Use 'zoi install' instead."
76 ));
77 }
78 };
79
80 println!(
81 "Currently installed version: {}",
82 old_manifest.version.yellow()
83 );
84 println!("Available version: {}", new_version.green());
85
86 if old_manifest.version == new_version {
87 println!("\nPackage is already up to date.");
88 return Ok(());
89 }
90
91 let download_size = new_pkg.archive_size.unwrap_or(0);
92 let old_installed_size = old_manifest.installed_size.unwrap_or(0);
93 let new_installed_size = new_pkg.installed_size.unwrap_or(0);
94 let installed_size_diff = new_installed_size as i64 - old_installed_size as i64;
95
96 println!();
97 println!(
98 "Total Download Size: {}",
99 crate::utils::format_bytes(download_size)
100 );
101 println!(
102 "Net Upgrade Size: {}",
103 crate::utils::format_size_diff(installed_size_diff)
104 );
105 println!();
106
107 if !crate::utils::ask_for_confirmation(
108 &format!("Update from {} to {}?", old_manifest.version, new_version),
109 yes,
110 ) {
111 return Ok(());
112 }
113
114 let transaction = transaction::begin()?;
115
116 if let Some(hooks) = &new_pkg.hooks {
117 hooks::run_hooks(hooks, hooks::HookType::PreUpgrade)?;
118 }
119
120 let (graph, _) = install::resolver::resolve_dependency_graph(
121 &[package_name.to_string()],
122 Some(old_manifest.scope),
123 true,
124 yes,
125 false,
126 None,
127 true,
128 )?;
129
130 let install_plan = install::plan::create_install_plan(&graph.nodes)?;
131
132 let mut new_manifest_option: Option<types::InstallManifest> = None;
133
134 for (id, node) in &graph.nodes {
135 if let Some(action) = install_plan.get(id) {
136 match install::installer::install_node(node, action, None, None, yes) {
137 Ok(m) => {
138 if m.name == new_pkg.name {
139 new_manifest_option = Some(m);
140 }
141 }
142 Err(e) => {
143 eprintln!("\nError: Update failed during installation. Rolling back...");
144 transaction::rollback(&transaction.id)?;
145 return Err(anyhow!("Update failed: {}", e));
146 }
147 }
148 }
149 }
150
151 if let Some(new_manifest) = new_manifest_option {
152 if let Err(e) = transaction::record_operation(
153 &transaction.id,
154 types::TransactionOperation::Upgrade {
155 old_manifest: Box::new(old_manifest.clone()),
156 new_manifest: Box::new(new_manifest.clone()),
157 },
158 ) {
159 eprintln!("Warning: Failed to record transaction for update: {}", e);
160 transaction::delete_log(&transaction.id)?;
161 } else {
162 transaction::commit(&transaction.id)?;
163 }
164
165 if let Some(backup_files) = &old_manifest.backup {
166 println!("Restoring configuration files...");
167 let old_version_dir = local::get_package_version_dir(
168 old_manifest.scope,
169 &old_manifest.registry_handle,
170 &old_manifest.repo,
171 &old_manifest.name,
172 &old_manifest.version,
173 )?;
174 let new_version_dir = local::get_package_version_dir(
175 new_manifest.scope,
176 &new_manifest.registry_handle,
177 &new_manifest.repo,
178 &new_manifest.name,
179 &new_manifest.version,
180 )?;
181
182 for backup_file_rel in backup_files {
183 let old_path = old_version_dir.join(backup_file_rel);
184 let new_path = new_version_dir.join(backup_file_rel);
185
186 if old_path.exists() {
187 if new_path.exists() {
188 let zoinew_path = new_path.with_extension(format!(
189 "{}.zoinew",
190 new_path
191 .extension()
192 .and_then(|s| s.to_str())
193 .unwrap_or_default()
194 ));
195 println!(
196 "Configuration file '{}' exists in new version. Saving as .zoinew",
197 new_path.display()
198 );
199 if let Err(e) = fs::rename(&new_path, &zoinew_path) {
200 eprintln!("Warning: failed to rename to .zoinew: {}", e);
201 continue;
202 }
203 }
204 if let Some(p) = new_path.parent() {
205 fs::create_dir_all(p)?;
206 }
207 if let Err(e) = fs::rename(&old_path, &new_path) {
208 eprintln!("Warning: failed to restore backup file: {}", e);
209 }
210 }
211 }
212 }
213
214 cleanup_old_versions(
215 &new_pkg.name,
216 old_manifest.scope,
217 &new_pkg.repo,
218 registry_handle.as_deref().unwrap_or("local"),
219 )?;
220
221 if let Some(hooks) = &new_pkg.hooks {
222 hooks::run_hooks(hooks, hooks::HookType::PostUpgrade)?;
223 }
224
225 println!("\n{}", "Success:".green());
226 Ok(())
227 } else {
228 eprintln!("\nError: Update failed to produce a new manifest. Rolling back...");
229 transaction::rollback(&transaction.id)?;
230 Err(anyhow!("Update failed: could not get new manifest"))
231 }
232}
233
234fn run_update_all_logic(yes: bool) -> Result<()> {
235 let installed_packages = local::get_installed_packages()?;
236 let pinned_packages = pin::get_pinned_packages()?;
237 let pinned_sources: Vec<String> = pinned_packages.into_iter().map(|p| p.source).collect();
238
239 let mut packages_to_upgrade = Vec::new();
240 let mut upgrade_messages = Vec::new();
241
242 println!("\n{}", "--- Checking for Upgrades ---".yellow().bold());
243 let pb = ProgressBar::new(installed_packages.len() as u64);
244 pb.set_style(
245 ProgressStyle::default_bar()
246 .template(
247 "{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} ({msg})",
248 )?
249 .progress_chars("#>-"),
250 );
251 pb.set_message("Checking packages...");
252
253 for manifest in installed_packages {
254 let source = format!("#{}@{}", manifest.registry_handle, manifest.repo);
255 if pinned_sources.contains(&source) {
256 upgrade_messages.push(format!("- {} is pinned, skipping.", manifest.name.cyan()));
257 continue;
258 }
259
260 let (new_pkg, new_version, _, _, registry_handle) =
261 match resolve::resolve_package_and_version(&source, true) {
262 Ok(result) => result,
263 Err(e) => {
264 upgrade_messages.push(format!(
265 "- Could not resolve package '{}': {}, skipping.",
266 manifest.name, e
267 ));
268 continue;
269 }
270 };
271
272 if manifest.version != new_version {
273 upgrade_messages.push(format!(
274 "- {} can be upgraded from {} to {}",
275 manifest.name.cyan(),
276 manifest.version.yellow(),
277 new_version.green()
278 ));
279 packages_to_upgrade.push((source.clone(), new_pkg, registry_handle, manifest));
280 } else {
281 upgrade_messages.push(format!("- {} is up to date.", manifest.name.cyan()));
282 }
283 pb.inc(1);
284 }
285 pb.finish_and_clear();
286
287 for msg in upgrade_messages {
288 println!("{}", msg);
289 }
290
291 if packages_to_upgrade.is_empty() {
292 println!("\nAll packages are up to date.");
293 return Ok(());
294 }
295
296 let total_download_size: u64 = packages_to_upgrade
297 .iter()
298 .map(|(_, pkg, _, _)| pkg.archive_size.unwrap_or(0))
299 .sum();
300
301 let total_installed_size_diff: i64 = packages_to_upgrade
302 .iter()
303 .map(|(_, new_pkg, _, old_manifest)| {
304 let old_size = old_manifest.installed_size.unwrap_or(0) as i64;
305 let new_size = new_pkg.installed_size.unwrap_or(0) as i64;
306 new_size - old_size
307 })
308 .sum();
309
310 println!();
311 println!(
312 "Total Download Size: {}",
313 crate::utils::format_bytes(total_download_size)
314 );
315 println!(
316 "Net Upgrade Size: {}",
317 crate::utils::format_size_diff(total_installed_size_diff)
318 );
319
320 println!();
321 if !crate::utils::ask_for_confirmation("Do you want to upgrade these packages?", yes) {
322 return Ok(());
323 }
324
325 let transaction = transaction::begin()?;
326 let failed_updates = Mutex::new(Vec::new());
327 let successful_upgrades = Mutex::new(Vec::new());
328
329 packages_to_upgrade
330 .par_iter()
331 .for_each(|(source, new_pkg, _registry_handle, old_manifest)| {
332 println!(
333 "\n--- Upgrading {} to {} ---",
334 source.cyan(),
335 new_pkg.version.as_deref().unwrap_or("N/A").green()
336 );
337
338 if let Some(hooks) = &new_pkg.hooks
339 && let Err(e) = hooks::run_hooks(hooks, hooks::HookType::PreUpgrade)
340 {
341 eprintln!(
342 "{}: Pre-upgrade hook failed for '{}': {}",
343 "Error".red().bold(),
344 source,
345 e
346 );
347 failed_updates.lock().unwrap().push(source.clone());
348 return;
349 }
350
351 let (graph, _) = match install::resolver::resolve_dependency_graph(
352 &[source.to_string()],
353 Some(old_manifest.scope),
354 true,
355 yes,
356 false,
357 None,
358 true,
359 ) {
360 Ok(res) => res,
361 Err(e) => {
362 eprintln!("Error resolving dependency graph for update: {}", e);
363 failed_updates.lock().unwrap().push(source.clone());
364 return;
365 }
366 };
367
368 let install_plan = match install::plan::create_install_plan(&graph.nodes) {
369 Ok(plan) => plan,
370 Err(e) => {
371 eprintln!("Error creating install plan for update: {}", e);
372 failed_updates.lock().unwrap().push(source.clone());
373 return;
374 }
375 };
376
377 let mut new_manifest_option: Option<types::InstallManifest> = None;
378
379 for (id, node) in &graph.nodes {
380 if let Some(action) = install_plan.get(id) {
381 match install::installer::install_node(node, action, None, None, yes) {
382 Ok(m) => {
383 if m.name == new_pkg.name {
384 new_manifest_option = Some(m);
385 }
386 }
387 Err(e) => {
388 eprintln!("Failed to upgrade {}: {}", source, e);
389 failed_updates.lock().unwrap().push(source.clone());
390 return;
391 }
392 }
393 }
394 }
395
396 if let Some(new_manifest) = new_manifest_option {
397 if let Err(e) = transaction::record_operation(
398 &transaction.id,
399 types::TransactionOperation::Upgrade {
400 old_manifest: Box::new(old_manifest.clone()),
401 new_manifest: Box::new(new_manifest.clone()),
402 },
403 ) {
404 eprintln!("Error: Failed to record transaction for {}: {}", source, e);
405 failed_updates.lock().unwrap().push(source.clone());
406 } else {
407 successful_upgrades.lock().unwrap().push((
408 old_manifest.clone(),
409 new_manifest.clone(),
410 new_pkg.clone(),
411 ));
412 }
413 } else {
414 eprintln!("Failed to get new manifest for {}", source);
415 failed_updates.lock().unwrap().push(source.clone());
416 }
417 });
418
419 let failed = failed_updates.into_inner().unwrap();
420 if !failed.is_empty() {
421 eprintln!("\nError: Some packages failed to upgrade. Rolling back all changes...");
422 for pkg in &failed {
423 eprintln!(" - {}", pkg);
424 }
425 transaction::rollback(&transaction.id)?;
426 return Err(anyhow!("Update failed for some packages."));
427 }
428
429 transaction::commit(&transaction.id)?;
430
431 println!("\n{}", "Success:".green());
432 let successful_upgrades = successful_upgrades.into_inner().unwrap();
433 for (old_manifest, new_manifest, new_pkg) in &successful_upgrades {
434 if let Some(backup_files) = &old_manifest.backup {
435 println!(
436 "Restoring configuration for {}...",
437 old_manifest.name.cyan()
438 );
439 let old_version_dir = local::get_package_version_dir(
440 old_manifest.scope,
441 &old_manifest.registry_handle,
442 &old_manifest.repo,
443 &old_manifest.name,
444 &old_manifest.version,
445 )?;
446 let new_version_dir = local::get_package_version_dir(
447 new_manifest.scope,
448 &new_manifest.registry_handle,
449 &new_manifest.repo,
450 &new_manifest.name,
451 &new_manifest.version,
452 )?;
453
454 for backup_file_rel in backup_files {
455 let old_path = old_version_dir.join(backup_file_rel);
456 let new_path = new_version_dir.join(backup_file_rel);
457
458 if old_path.exists() {
459 if new_path.exists() {
460 let zoinew_path = new_path.with_extension(format!(
461 "{}.zoinew",
462 new_path
463 .extension()
464 .and_then(|s| s.to_str())
465 .unwrap_or_default()
466 ));
467 println!(
468 "Configuration file '{}' exists in new version. Saving as .zoinew",
469 new_path.display()
470 );
471 if let Err(e) = fs::rename(&new_path, &zoinew_path) {
472 eprintln!("Warning: failed to rename to .zoinew: {}", e);
473 continue;
474 }
475 }
476 if let Some(p) = new_path.parent() {
477 fs::create_dir_all(p)?;
478 }
479 if let Err(e) = fs::rename(&old_path, &new_path) {
480 eprintln!("Warning: failed to restore backup file: {}", e);
481 }
482 }
483 }
484 }
485
486 if let Err(e) = cleanup_old_versions(
487 &new_manifest.name,
488 new_manifest.scope,
489 &new_manifest.repo,
490 &new_manifest.registry_handle,
491 ) {
492 eprintln!(
493 "Failed to clean up old versions for {}: {}",
494 new_manifest.name, e
495 );
496 }
497
498 if let Some(hooks) = &new_pkg.hooks
499 && let Err(e) = hooks::run_hooks(hooks, hooks::HookType::PostUpgrade)
500 {
501 eprintln!(
502 "{}: Post-upgrade hook failed for '{}': {}",
503 "Error".red().bold(),
504 new_manifest.name,
505 e
506 );
507 }
508 }
509
510 println!("\n{}", "Success:".green());
511 Ok(())
512}
513
514fn cleanup_old_versions(
515 package_name: &str,
516 scope: types::Scope,
517 repo: &str,
518 registry_handle: &str,
519) -> Result<()> {
520 let config = config::read_config()?;
521 let rollback_enabled = config.rollback_enabled;
522
523 let package_dir = local::get_package_dir(scope, registry_handle, repo, package_name)?;
524
525 let mut versions = Vec::new();
526 if let Ok(entries) = fs::read_dir(&package_dir) {
527 for entry in entries.flatten() {
528 let path = entry.path();
529 if path.is_dir()
530 && let Some(version_str) = path.file_name().and_then(|s| s.to_str())
531 && version_str != "latest"
532 && let Ok(version) = Version::parse(version_str)
533 {
534 versions.push(version);
535 }
536 }
537 }
538
539 if versions.is_empty() {
540 return Ok(());
541 }
542
543 versions.sort();
544
545 let versions_to_keep = if rollback_enabled { 2 } else { 1 };
546
547 if versions.len() > versions_to_keep {
548 let num_to_delete = versions.len() - versions_to_keep;
549 let versions_to_delete = &versions[..num_to_delete];
550
551 println!("Cleaning up old versions...");
552 for version in versions_to_delete {
553 let version_dir_to_delete = package_dir.join(version.to_string());
554 println!(" - Removing {}", version_dir_to_delete.display());
555 if version_dir_to_delete.exists() {
556 fs::remove_dir_all(version_dir_to_delete)?;
557 }
558 }
559 }
560
561 Ok(())
562}