From 4596716cc8f23353176c61b190d38f552b2b33a4 Mon Sep 17 00:00:00 2001
From: Cole Helbling <cole.e.helbling@outlook.com>
Date: Thu, 12 Mar 2020 16:45:06 -0700
Subject: [PATCH] Migrate from `z` database (#29)

Migrate from `z` database

The new `migrate` subcommand takes in a path to the old `z` database and
naively parses it to add to the database. The command will fail if the user already has a database, so as to prevent tainting it in any way.
---
 src/db.rs   | 90 ++++++++++++++++++++++++++++++++++++++++++++++++++---
 src/main.rs |  8 +++++
 2 files changed, 94 insertions(+), 4 deletions(-)

diff --git a/src/db.rs b/src/db.rs
index cf9130d..6b951ba 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -1,10 +1,10 @@
 use crate::dir::Dir;
 use crate::types::{Rank, Timestamp};
-use crate::util::get_zo_maxage;
-use anyhow::{anyhow, Context, Result};
+use crate::util;
+use anyhow::{anyhow, bail, Context, Result};
 use fs2::FileExt;
 use std::fs::{self, File, OpenOptions};
-use std::io::{self, BufReader, BufWriter};
+use std::io::{self, BufRead, BufReader, BufWriter};
 use std::path::{Path, PathBuf};
 
 pub struct DB {
@@ -71,6 +71,88 @@ impl DB {
         Ok(())
     }
 
+    pub fn migrate<P: AsRef<Path>>(&mut self, path: P) -> Result<()> {
+        if !self.dirs.is_empty() {
+            bail!(
+                "To prevent conflicts, you can only migrate from z with an empty \
+                 zoxide database!"
+            );
+        }
+
+        let z_db_file =
+            File::open(path).with_context(|| anyhow!("could not open z database file"))?;
+        let reader = BufReader::new(z_db_file);
+
+        for (idx, read_line) in reader.lines().enumerate() {
+            let line_number = idx + 1;
+            let line = if let Ok(line) = read_line {
+                line
+            } else {
+                eprintln!("could not read line {}: {:?}", line_number, read_line);
+                continue;
+            };
+
+            let split_line = line.rsplitn(3, '|').collect::<Vec<&str>>();
+
+            match split_line.as_slice() {
+                [epoch_str, rank_str, path_str] => {
+                    let epoch = match epoch_str.parse::<i64>() {
+                        Ok(epoch) => epoch,
+                        Err(e) => {
+                            eprintln!(
+                                "invalid epoch '{}' at line {}: {}",
+                                epoch_str, line_number, e
+                            );
+                            continue;
+                        }
+                    };
+                    let rank = match rank_str.parse::<f64>() {
+                        Ok(rank) => rank,
+                        Err(e) => {
+                            eprintln!("invalid rank '{}' at line {}: {}", rank_str, line_number, e);
+                            continue;
+                        }
+                    };
+                    let path_abs = match Path::new(path_str).canonicalize() {
+                        Ok(path) => path,
+                        Err(e) => {
+                            eprintln!("invalid path '{}' at line {}: {}", path_str, line_number, e);
+                            continue;
+                        }
+                    };
+                    let path_str = match path_abs.to_str() {
+                        Some(path) => path,
+                        None => {
+                            eprintln!(
+                                "invalid unicode in path '{}' at line {}",
+                                path_abs.display(),
+                                line_number
+                            );
+                            continue;
+                        }
+                    };
+
+                    // FIXME: When we switch to PathBuf for storing directories inside Dir, just
+                    // pass `PathBuf::from(path_str)`
+                    self.dirs.push(Dir {
+                        path: path_str.to_string(),
+                        rank,
+                        last_accessed: epoch,
+                    });
+                }
+                [] | [""] => {} // ignore blank lines
+                line => {
+                    eprintln!("invalid entry at line {}: {:?}", line_number, line);
+                    continue;
+                }
+            };
+        }
+
+        self.modified = true;
+
+        Ok(())
+    }
+
     pub fn add<P: AsRef<Path>>(&mut self, path: P, now: Timestamp) -> Result<()> {
         let path_abs = path
             .as_ref()
@@ -93,7 +175,7 @@ impl DB {
             }
         };
 
-        let max_age = get_zo_maxage()?;
+        let max_age = util::get_zo_maxage()?;
         let sum_age = self.dirs.iter().map(|dir| dir.rank).sum::<Rank>();
 
         if sum_age > max_age {
diff --git a/src/main.rs b/src/main.rs
index 311aeb7..7a60442 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -28,6 +28,9 @@ enum Zoxide {
     #[structopt(about = "Add a new directory or increment its rank")]
     Add { path: Option<String> },
 
+    #[structopt(about = "Migrate from z database")]
+    Migrate { path: String },
+
     #[structopt(about = "Prints shell configuration")]
     Init {
         #[structopt(possible_values = &Shell::variants(), case_insensitive = true)]
@@ -94,6 +97,10 @@ pub fn main() -> Result<()> {
                 }
             }?;
         }
+        Zoxide::Migrate { path } => {
+            let mut db = get_db()?;
+            db.migrate(path)?;
+        }
         Zoxide::Init {
             shell,
             no_define_aliases,
@@ -208,6 +215,7 @@ end
 
 const INIT_FISH_ALIAS: &str = r#"
 abbr -a zi 'z -i'
+
 abbr -a za 'zoxide add'
 abbr -a zq 'zoxide query'
 abbr -a zr 'zoxide remove'