affine_importer 0.1.2

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

use nanoid::nanoid;

use super::{
  ImportAssetRef, ImportBatch, ImportBatchLimits, ImportCursor, ImportDocRef, ImportError, ImportOptions,
  ImportProgress, ImportResult, ImportedDocMeta, ImportedDocSnapshot, ImportedIcon, IndexedArchive, add_asset_batch,
  doc_snapshot, empty_batch, extract_html_icon, extract_html_title, file_name, has_emitted_assets, hash_bytes,
  html_to_markdown, normalize_import_path, push_skipped_doc_warning, register_page_path, strip_extension,
  strip_notion_hash,
};

pub(super) struct NotionHtmlImportCursor {
  archive: IndexedArchive,
  docs: Vec<ImportDocRef>,
  assets: Vec<ImportAssetRef>,
  page_ids_by_path: BTreeMap<String, String>,
  blob_ids_by_path: BTreeMap<String, String>,
  folders: Vec<crate::FolderHierarchyDelta>,
  options: ImportOptions,
  limits: ImportBatchLimits,
  next_doc: usize,
  emitted_assets: BTreeSet<usize>,
  emitted_final: bool,
  completed: usize,
  total: usize,
  entry_id: Option<String>,
  is_workspace_file: bool,
}

impl NotionHtmlImportCursor {
  pub(super) fn new(archive: IndexedArchive, options: ImportOptions, limits: ImportBatchLimits) -> ImportResult<Self> {
    let mut docs = Vec::new();
    let mut assets = Vec::new();
    let mut page_ids_by_path = BTreeMap::new();
    let mut blob_ids_by_path = BTreeMap::new();
    let mut is_workspace_file = false;
    for (entry_index, entry) in archive.entries.iter().enumerate() {
      let lower = entry.meta.path.to_lowercase();
      if lower.ends_with("/index.html") {
        is_workspace_file = true;
        continue;
      }
      if lower.ends_with(".html") {
        let doc_id = nanoid!();
        register_page_path(&mut page_ids_by_path, &entry.meta.path, &doc_id);
        let html = archive.read_entry(entry)?;
        let html = String::from_utf8_lossy(&html);
        docs.push(ImportDocRef {
          entry_index,
          doc_id,
          title: extract_html_title(&html),
          icon: extract_html_icon(&html),
        });
      }
    }
    for (entry_index, entry) in archive.entries.iter().enumerate() {
      let lower = entry.meta.path.to_lowercase();
      if !lower.ends_with(".html") && !lower.ends_with(".md") {
        let blob_id = hash_bytes(&archive.read_entry(entry)?);
        blob_ids_by_path.insert(entry.meta.path.clone(), blob_id.clone());
        assets.push(ImportAssetRef { entry_index, blob_id });
      }
    }
    if docs.is_empty() {
      return Err(ImportError::InvalidSource(
        "No Notion HTML pages found in the archive".to_string(),
      ));
    }
    let folders = notion_html_cursor_folders(&docs, &archive);
    let total = docs.len();
    let entry_id = docs.first().map(|doc| doc.doc_id.clone());
    Ok(Self {
      archive,
      docs,
      assets,
      page_ids_by_path,
      blob_ids_by_path,
      folders,
      options,
      limits,
      next_doc: 0,
      emitted_assets: BTreeSet::new(),
      emitted_final: false,
      completed: 0,
      total,
      entry_id,
      is_workspace_file,
    })
  }
}

impl ImportCursor for NotionHtmlImportCursor {
  fn total(&self) -> usize {
    self.total
  }

  fn next_batch(&mut self) -> ImportResult<Option<ImportBatch>> {
    if self.completed >= self.total && self.emitted_final {
      return Ok(None);
    }
    if self.options.cancel_after_entries == Some(self.completed) {
      return Err(ImportError::Cancelled);
    }

    let mut batch = empty_batch(self.total);
    batch.entry_id = self.entry_id.clone();
    batch.is_workspace_file = self.is_workspace_file;

    let end = (self.next_doc + self.limits.max_docs.max(1)).min(self.docs.len());
    for doc in &self.docs[self.next_doc..end] {
      let entry = &self.archive.entries[doc.entry_index];
      let bytes = self.archive.read_entry(entry)?;
      let html = String::from_utf8_lossy(&bytes);
      let markdown = html_to_markdown(&html);
      let title = doc
        .title
        .clone()
        .unwrap_or_else(|| strip_extension(file_name(&entry.meta.path)).to_string());
      let snapshot = match doc_snapshot(
        &doc.doc_id,
        &title,
        &markdown,
        &entry.meta.path,
        &self.page_ids_by_path,
        &self.blob_ids_by_path,
      ) {
        Ok(snapshot) => snapshot,
        Err(ImportError::Document(error)) => {
          push_skipped_doc_warning(&mut batch, &entry.meta.path, error);
          continue;
        }
        Err(error) => return Err(error),
      };
      batch.docs.push(ImportedDocSnapshot {
        id: doc.doc_id.clone(),
        snapshot,
        meta: Some(ImportedDocMeta {
          title: Some(title),
          create_date: None,
          updated_date: None,
          tags: None,
          favorite: None,
          trash: None,
        }),
      });
      if let Some(icon) = doc.icon.clone() {
        batch.icons.push(ImportedIcon {
          doc_id: doc.doc_id.clone(),
          icon,
        });
      }
    }
    add_asset_batch(
      &self.archive,
      &self.limits,
      &mut self.emitted_assets,
      &mut batch,
      &self.assets,
    )?;
    if end == self.docs.len() && has_emitted_assets(&self.emitted_assets, &self.assets) {
      batch.folders.extend(self.folders.clone());
      self.emitted_final = true;
    }
    self.next_doc = end;
    self.completed = end;
    batch.progress = ImportProgress {
      completed: self.completed,
      total: self.total,
    };
    batch.done = self.completed >= self.total && self.emitted_final;
    Ok(Some(batch))
  }
}

fn notion_html_cursor_folders(docs: &[ImportDocRef], archive: &IndexedArchive) -> Vec<crate::FolderHierarchyDelta> {
  let doc_paths = docs
    .iter()
    .map(|doc| {
      (
        archive.entries[doc.entry_index].meta.path.clone(),
        doc.doc_id.clone(),
        doc.icon.clone(),
      )
    })
    .collect::<Vec<_>>();
  notion_html_folders(&doc_paths)
}

fn notion_html_folders(
  doc_paths: &[(String, String, Option<crate::ImportedIconData>)],
) -> Vec<crate::FolderHierarchyDelta> {
  let mut page_by_folder = BTreeMap::<String, (String, Option<crate::ImportedIconData>)>::new();
  for (path, doc_id, icon) in doc_paths {
    page_by_folder.insert(
      strip_extension(&normalize_import_path(path)).to_string(),
      (doc_id.clone(), icon.clone()),
    );
  }

  let mut folders = BTreeMap::<String, crate::FolderHierarchyDelta>::new();
  for (path, doc_id, icon) in doc_paths {
    let normalized = normalize_import_path(path);
    let mut parts = normalized.split('/').map(ToString::to_string).collect::<Vec<_>>();
    let Some(file_name) = parts.pop() else {
      continue;
    };
    let leaf_name = strip_notion_hash(strip_extension(&file_name));
    let mut current_path = String::new();
    for folder in parts {
      let parent_path = (!current_path.is_empty()).then(|| current_path.clone());
      current_path = if current_path.is_empty() {
        folder.clone()
      } else {
        format!("{current_path}/{folder}")
      };
      let page = page_by_folder.get(&current_path);
      folders
        .entry(current_path.clone())
        .or_insert_with(|| crate::FolderHierarchyDelta {
          path: current_path.clone(),
          name: strip_notion_hash(&folder),
          parent_path,
          page_id: page.map(|(page_id, _)| page_id.clone()),
          icon: page.and_then(|(_, icon)| icon.clone()),
        });
    }
    let parent_path = (!current_path.is_empty()).then(|| current_path.clone());
    current_path = if current_path.is_empty() {
      leaf_name.clone()
    } else {
      format!("{current_path}/{leaf_name}")
    };
    folders.insert(
      current_path.clone(),
      crate::FolderHierarchyDelta {
        path: current_path,
        name: leaf_name,
        parent_path,
        page_id: Some(doc_id.clone()),
        icon: icon.clone(),
      },
    );
  }
  folders.into_values().collect()
}