affine_importer 0.1.2

AFFiNE import planning and batching core.
Documentation
use std::collections::{BTreeMap, BTreeSet};

use sha2::{Digest, Sha256};

mod frontmatter;
mod html;
mod providers;
mod references;
mod utils;
use frontmatter::{frontmatter_meta, parse_date_millis, strip_frontmatter};
use html::{extract_html_icon, extract_html_title, html_to_markdown, split_leading_emoji};
pub(crate) use providers::register_builtin_providers;
use references::{
  ImportedAssetMetadata, imported_asset_metadata, is_protected, linked_page_reference, parse_link_span,
  register_csv_path, register_page_path, resolve_path, rewrite_asset_references, rewrite_bare_page_references,
  rewrite_link_references, rewrite_markdown_lines, rewrite_markdown_link_destinations,
  rewrite_non_image_embeds_to_attachments, rewrite_page_references,
};
use utils::{file_name, mime_from_path, strip_extension};

use crate::{
  ArchiveEntryMeta, ArchiveSource, DirectoryPathSource, FolderHierarchyDelta, ImportBatch, ImportError, ImportProgress,
  ImportResult, ImportSource, ImportWarning, ImportedAsset, ImportedIconData, ZipPathSource,
  vfs::{VfsEntry, normalize_import_path},
};

#[derive(Debug, Clone, Default)]
pub struct ImportOptions {
  pub cancel: bool,
  pub cancel_after_entries: Option<usize>,
}

#[derive(Debug, Clone)]
pub struct ImportBatchLimits {
  pub max_docs: usize,
  pub max_blobs: usize,
  pub max_blob_bytes: u64,
}

impl Default for ImportBatchLimits {
  fn default() -> Self {
    Self {
      max_docs: 20,
      max_blobs: 100,
      max_blob_bytes: 10 * 1024 * 1024,
    }
  }
}

#[derive(Debug, Clone)]
struct IndexedEntry {
  meta: ArchiveEntryMeta,
  bytes: Option<Vec<u8>>,
}

struct IndexedArchive {
  source: Box<dyn ArchiveSource>,
  entries: Vec<IndexedEntry>,
}

impl IndexedArchive {
  fn read<S: ArchiveSource + 'static>(source: S, expand_nested_zips: bool) -> ImportResult<Self> {
    let mut archive = Self {
      entries: source
        .entries()?
        .into_iter()
        .map(|meta| IndexedEntry { meta, bytes: None })
        .collect(),
      source: Box::new(source),
    };
    if expand_nested_zips {
      archive.expand_nested_zips()?;
    }
    Ok(archive)
  }

  fn expand_nested_zips(&mut self) -> ImportResult<()> {
    let mut expanded = Vec::new();
    for entry in self.entries.drain(..).collect::<Vec<_>>() {
      if entry.meta.path.to_lowercase().ends_with(".zip") {
        let base_path = strip_extension(&entry.meta.path).to_string();
        let bytes = if let Some(bytes) = entry.bytes {
          bytes
        } else {
          self.source.read_entry(entry.meta.index)?
        };
        let nested = read_zip_bytes_entries(&bytes, &base_path)?;
        expanded.extend(nested);
      } else {
        expanded.push(entry);
      }
    }
    self.entries = expanded;
    Ok(())
  }

  fn read_entry(&self, entry: &IndexedEntry) -> ImportResult<Vec<u8>> {
    if let Some(bytes) = &entry.bytes {
      return Ok(bytes.clone());
    }
    self.source.read_entry(entry.meta.index)
  }
}

fn read_zip_bytes_entries(bytes: &[u8], base_path: &str) -> ImportResult<Vec<IndexedEntry>> {
  use std::io::{Cursor, Read};

  let mut archive = zip::ZipArchive::new(Cursor::new(bytes))?;
  let mut entries = Vec::new();
  for index in 0..archive.len() {
    let mut file = archive.by_index(index)?;
    if file.is_dir() {
      continue;
    }
    let Some(path) = file
      .enclosed_name()
      .map(|path| normalize_import_path(&path.to_string_lossy()))
    else {
      continue;
    };
    if path.is_empty() || crate::source::is_system_path(&path) {
      continue;
    }
    let mut bytes = Vec::with_capacity(file.size() as usize);
    file.read_to_end(&mut bytes)?;
    let path = if base_path.is_empty() {
      path
    } else {
      format!("{base_path}/{path}")
    };
    entries.push(IndexedEntry {
      meta: ArchiveEntryMeta {
        index,
        path,
        compressed_size: file.compressed_size(),
        uncompressed_size: file.size(),
      },
      bytes: Some(bytes),
    });
  }
  Ok(entries)
}

#[derive(Debug, Clone)]
struct ImportAssetRef {
  entry_index: usize,
  blob_id: String,
}

#[derive(Debug, Clone)]
struct ImportDocRef {
  entry_index: usize,
  doc_id: String,
  title: Option<String>,
  icon: Option<ImportedIconData>,
}

fn read_archive(
  source: ImportSource,
  expand_nested_zips: bool,
  options: &ImportOptions,
) -> ImportResult<IndexedArchive> {
  if options.cancel {
    return Err(ImportError::Cancelled);
  }
  match source {
    ImportSource::FilePath(path) => IndexedArchive::read(ZipPathSource::new(path), expand_nested_zips),
    ImportSource::DirectoryPath(path) => IndexedArchive::read(DirectoryPathSource::new(path), expand_nested_zips),
  }
}

fn empty_batch(total: usize) -> ImportBatch {
  ImportBatch {
    docs: Vec::new(),
    blobs: Vec::new(),
    folders: Vec::new(),
    tags: Vec::new(),
    icons: Vec::new(),
    warnings: Vec::new(),
    progress: ImportProgress { completed: 0, total },
    entry_id: None,
    is_workspace_file: false,
    done: false,
  }
}

fn hash_bytes(bytes: &[u8]) -> String {
  let mut hasher = Sha256::new();
  hasher.update(bytes);
  hasher
    .finalize()
    .iter()
    .map(|byte| format!("{byte:02x}"))
    .collect::<String>()
}

fn entry_from_indexed(entry: &IndexedEntry, bytes: Vec<u8>) -> VfsEntry {
  VfsEntry {
    path: entry.meta.path.clone(),
    bytes,
  }
}

fn asset_from_indexed(
  archive: &IndexedArchive,
  entry: &IndexedEntry,
  asset: &ImportAssetRef,
) -> ImportResult<ImportedAsset> {
  Ok(ImportedAsset {
    blob_id: asset.blob_id.clone(),
    source_path: entry.meta.path.clone(),
    file_name: file_name(&entry.meta.path).to_string(),
    mime: mime_from_path(&entry.meta.path).to_string(),
    bytes: archive.read_entry(entry)?,
  })
}

fn add_asset_batch(
  archive: &IndexedArchive,
  limits: &ImportBatchLimits,
  emitted_assets: &mut BTreeSet<usize>,
  batch: &mut ImportBatch,
  assets: &[ImportAssetRef],
) -> ImportResult<()> {
  let mut bytes_in_batch = 0u64;
  let max_blobs = limits.max_blobs.max(1);
  for (asset_index, asset) in assets.iter().enumerate() {
    if emitted_assets.contains(&asset_index) || batch.blobs.len() >= max_blobs {
      continue;
    }
    let entry = &archive.entries[asset.entry_index];
    if entry.meta.uncompressed_size > limits.max_blob_bytes {
      batch.warnings.push(ImportWarning {
        code: "skipped_asset".to_string(),
        source_path: Some(entry.meta.path.clone()),
        message: format!(
          "Skipped {}: asset is larger than the batch import support",
          entry.meta.path
        ),
      });
      emitted_assets.insert(asset_index);
      continue;
    }
    if !batch.blobs.is_empty() && bytes_in_batch + entry.meta.uncompressed_size > limits.max_blob_bytes {
      continue;
    }
    let imported = asset_from_indexed(archive, entry, asset)?;
    bytes_in_batch += entry.meta.uncompressed_size;
    batch.blobs.push(imported);
    emitted_assets.insert(asset_index);
  }
  Ok(())
}

fn has_emitted_assets(emitted_assets: &BTreeSet<usize>, assets: &[ImportAssetRef]) -> bool {
  emitted_assets.len() >= assets.len()
}

fn folders_for_doc_path(path: &str, doc_id: &str, icon: Option<ImportedIconData>) -> Vec<FolderHierarchyDelta> {
  folder_hierarchy_deltas(folder_parts(path), Some(doc_id), icon, |folder| folder.to_string())
}

fn folder_hierarchy_deltas<I, F>(
  parts: I,
  doc_id: Option<&str>,
  icon: Option<ImportedIconData>,
  mut name_for_part: F,
) -> Vec<FolderHierarchyDelta>
where
  I: IntoIterator,
  I::Item: AsRef<str>,
  F: FnMut(&str) -> String,
{
  let mut folders = Vec::new();
  let mut current_path = String::new();
  for part in parts {
    let folder = part.as_ref();
    let parent_path = (!current_path.is_empty()).then(|| current_path.clone());
    current_path = if current_path.is_empty() {
      folder.to_string()
    } else {
      format!("{current_path}/{folder}")
    };
    folders.push(FolderHierarchyDelta {
      path: current_path.clone(),
      name: name_for_part(folder),
      parent_path,
      page_id: None,
      icon: None,
    });
  }
  if let Some(doc_id) = doc_id
    && !current_path.is_empty()
  {
    folders.push(FolderHierarchyDelta {
      path: format!("{current_path}/__doc__{doc_id}"),
      name: format!("__doc__{doc_id}"),
      parent_path: Some(current_path),
      page_id: Some(doc_id.to_string()),
      icon,
    });
  }
  folders
}

fn doc_snapshot(
  doc_id: &str,
  title: &str,
  markdown: &str,
  current_path: &str,
  page_ids_by_path: &BTreeMap<String, String>,
  blob_ids_by_path: &BTreeMap<String, String>,
) -> ImportResult<serde_json::Value> {
  doc_snapshot_inner(
    doc_id,
    title,
    markdown,
    current_path,
    page_ids_by_path,
    blob_ids_by_path,
    None,
  )
}

fn doc_snapshot_with_id_hints(
  doc_id: &str,
  title: &str,
  markdown: &str,
  current_path: &str,
  page_ids_by_path: &BTreeMap<String, String>,
  blob_ids_by_path: &BTreeMap<String, String>,
  id_hints: (&str, &str),
) -> ImportResult<serde_json::Value> {
  doc_snapshot_inner(
    doc_id,
    title,
    markdown,
    current_path,
    page_ids_by_path,
    blob_ids_by_path,
    Some(id_hints),
  )
}

fn doc_snapshot_inner(
  doc_id: &str,
  title: &str,
  markdown: &str,
  current_path: &str,
  page_ids_by_path: &BTreeMap<String, String>,
  blob_ids_by_path: &BTreeMap<String, String>,
  id_hints: Option<(&str, &str)>,
) -> ImportResult<serde_json::Value> {
  let rewritten = rewrite_markdown_link_destinations(markdown, current_path, page_ids_by_path);
  let rewritten = rewrite_bare_page_references(&rewritten, current_path, page_ids_by_path);
  let rewritten = rewrite_asset_references(&rewritten, current_path, blob_ids_by_path);
  let mut snapshot = match id_hints {
    Some((namespace, token)) => {
      affine_doc_loader::build_doc_snapshot_with_id_hints(title, &rewritten, doc_id, namespace, token)?
    }
    None => affine_doc_loader::build_doc_snapshot(title, &rewritten, doc_id)?,
  };
  rewrite_page_references(&mut snapshot, current_path, page_ids_by_path);
  Ok(snapshot)
}

fn push_skipped_doc_warning(batch: &mut ImportBatch, source_path: &str, error: affine_doc_loader::ParseError) {
  batch.warnings.push(ImportWarning {
    code: "skipped_doc".to_string(),
    source_path: Some(source_path.to_string()),
    message: format!("Skipped {source_path}: {error}"),
  });
}

fn folder_parts(path: &str) -> Vec<String> {
  let mut parts = normalize_import_path(path)
    .split('/')
    .map(ToString::to_string)
    .collect::<Vec<_>>();
  parts.pop();
  if parts.len() > 1 {
    parts.remove(0);
  }
  parts
}