repo_description: Qdrant - High-performance, massive-scale Vector Database and Vector Search Engine for the next generation of AI. Also available in the cloud https://cloud.qdrant.io/
repo_topics: ai-search ai-search-engine embeddings-similarity hnsw hybrid-search image-search knn-algorithm machine-learning mlops nearest-neighbor-search neural-network neural-search recommender-system search search-engine search-engines similarity-search vector-database vector-search vector-search-engine
code: use std::io::BufReader;
  use std::path::{Path, PathBuf};
  
  use collection::collection::Collection;
  use collection::shards::shard::PeerId;
  use common::fs::safe_delete_in_tmp;
  use common::tar_unpack::tar_unpack_file;
  use fs_err as fs;
  use fs_err::File;
  use log::info;
  use shard::snapshots::snapshot_data::SnapshotData;
  use storage::content_manager::alias_mapping::AliasPersistence;
  use storage::content_manager::snapshots::SnapshotConfig;
  use storage::content_manager::toc::{ALIASES_PATH, COLLECTIONS_DIR};
  
  /// Recover snapshots from the given arguments
  ///
  /// # Arguments
  ///
  /// * `mapping` - `[ <path>:<collection_name> ]`
  /// * `force` - if true, allow to overwrite collections from snapshots
  ///
  /// # Returns
  ///
  /// * `Vec<String>` - list of collections that were recovered
  pub fn recover_snapshots(
      mapping: &[String],
      force: bool,
      temp_dir: Option<&Path>,
      storage_dir: &Path,
      this_peer_id: PeerId,
      is_distributed: bool,
  ) -> Vec<String> {
      let collection_dir_path = storage_dir.join(COLLECTIONS_DIR);
      let mut recovered_collections: Vec<String> = vec![];
  
      for snapshot_params in mapping {
          let mut split = snapshot_params.split(':');
          let path = split
              .next()
              .unwrap_or_else(|| panic!("Snapshot path is missing: {snapshot_params}"));
  
          let snapshot_data = SnapshotData::new_packed_persistent(path);
  
          let collection_name = split
              .next()
              .unwrap_or_else(|| panic!("Collection name is missing: {snapshot_params}"));
          recovered_collections.push(collection_name.to_string());
          assert!(
              split.next().is_none(),
              "Too many parts in snapshot mapping: {snapshot_params}"
          );
          info!("Recovering snapshot {collection_name} from {path}");
          // check if collection already exists
          // if it does, we need to check if we want to overwrite it
          // if not, we need to abort
          let collection_path = collection_dir_path.join(collection_name);
          info!("Collection path: {}", collection_path.display());
          if collection_path.exists() {
              if !force {
                  panic!(
                      "Collection {collection_name} already exists. Use --force-snapshot to overwrite it."
                  );
              }
              info!("Overwriting collection {collection_name}");
          }
          let collection_temp_path =
              temp_dir.map_or_else(|| collection_path.with_extension("tmp"), PathBuf::from);
          if let Err(err) = Collection::restore_snapshot(
              snapshot_data,
              &collection_temp_path,
              this_peer_id,
              is_distributed,
          ) {
              panic!("Failed to recover snapshot {collection_name}: {err}");
          }
          // Remove collection_path directory if exists
          if collection_path.exists()
              && let Err(err) = safe_delete_in_tmp(&collection_path, &storage_dir.join(".deleted"))
                  .and_then(|to_delete| to_delete.close())
          {
              panic!("Failed to remove collection {collection_name}: {err}");
          }
          fs::rename(&collection_temp_path, &collection_path).unwrap();
      }
      recovered_collections
  }
  
  pub fn recover_full_snapshot(
      temp_dir: Option<&Path>,
      snapshot_path: &str,
      storage_dir: &Path,
      force: bool,
      this_peer_id: PeerId,
      is_distributed: bool,
  ) -> Vec<String> {
      let snapshot_temp_path = temp_dir
          .map(PathBuf::from)
          .unwrap_or_else(|| storage_dir.join("snapshots_recovery_tmp"));
      fs::create_dir_all(&snapshot_temp_path).unwrap();
  
      // Un-tar snapshot into temporary directory
      tar_unpack_file(Path::new(snapshot_path), &snapshot_temp_path).unwrap();
  
      // Read configuration file with snapshot-to-collection mapping
      let config_path = snapshot_temp_path.join("config.json");
      let config_file = BufReader::new(File::open(config_path).unwrap());
      let config_json: SnapshotConfig = serde_json::from_reader(config_file).unwrap();
  
      // Create mapping from the configuration file
      let mapping: Vec<String> = config_json
          .collections_mapping
          .iter()
          .map(|(collection_name, snapshot_file)| {
              format!(
                  "{}:{collection_name}",
                  snapshot_temp_path.join(snapshot_file).to_str().unwrap(),
              )
          })
          .collect();
  
      // Launch regular recovery of snapshots
      let recovered_collection = recover_snapshots(
          &mapping,
          force,
          temp_dir,
          storage_dir,
          this_peer_id,
          is_distributed,
      );
  
      let alias_path = storage_dir.join(ALIASES_PATH);
      let mut alias_persistence =
          AliasPersistence::open(&alias_path).expect("Can't open database by the provided config");
      for (alias, collection_name) in config_json.collections_aliases {
          if alias_persistence.get(&alias).is_some() && !force {
              panic!("Alias {alias} already exists. Use --force-snapshot to overwrite it.");
          }
          alias_persistence.insert(alias, collection_name).unwrap();
      }
  
      // Remove temporary directory
      fs::remove_dir_all(&snapshot_temp_path).unwrap();
      recovered_collection
  }
  