1use std::collections::{BTreeMap, BTreeSet};
2
3use sha2::{Digest, Sha256};
4
5mod frontmatter;
6mod html;
7mod providers;
8mod references;
9mod utils;
10use frontmatter::{frontmatter_meta, parse_date_millis, strip_frontmatter};
11use html::{extract_html_icon, extract_html_title, html_to_markdown, split_leading_emoji};
12pub(crate) use providers::register_builtin_providers;
13use references::{
14 ImportedAssetMetadata, imported_asset_metadata, is_protected, linked_page_reference, parse_link_span,
15 register_csv_path, register_page_path, resolve_path, rewrite_asset_references, rewrite_bare_page_references,
16 rewrite_link_references, rewrite_markdown_lines, rewrite_markdown_link_destinations,
17 rewrite_non_image_embeds_to_attachments, rewrite_page_references,
18};
19use utils::{file_name, mime_from_path, strip_extension};
20
21use crate::{
22 ArchiveEntryMeta, ArchiveSource, DirectoryPathSource, FolderHierarchyDelta, ImportBatch, ImportError, ImportProgress,
23 ImportResult, ImportSource, ImportWarning, ImportedAsset, ImportedIconData, ZipPathSource,
24 vfs::{VfsEntry, normalize_import_path},
25};
26
27#[derive(Debug, Clone, Default)]
28pub struct ImportOptions {
29 pub cancel: bool,
30 pub cancel_after_entries: Option<usize>,
31}
32
33#[derive(Debug, Clone)]
34pub struct ImportBatchLimits {
35 pub max_docs: usize,
36 pub max_blobs: usize,
37 pub max_blob_bytes: u64,
38}
39
40impl Default for ImportBatchLimits {
41 fn default() -> Self {
42 Self {
43 max_docs: 20,
44 max_blobs: 100,
45 max_blob_bytes: 10 * 1024 * 1024,
46 }
47 }
48}
49
50#[derive(Debug, Clone)]
51struct IndexedEntry {
52 meta: ArchiveEntryMeta,
53 bytes: Option<Vec<u8>>,
54}
55
56struct IndexedArchive {
57 source: Box<dyn ArchiveSource>,
58 entries: Vec<IndexedEntry>,
59}
60
61impl IndexedArchive {
62 fn read<S: ArchiveSource + 'static>(source: S, expand_nested_zips: bool) -> ImportResult<Self> {
63 let mut archive = Self {
64 entries: source
65 .entries()?
66 .into_iter()
67 .map(|meta| IndexedEntry { meta, bytes: None })
68 .collect(),
69 source: Box::new(source),
70 };
71 if expand_nested_zips {
72 archive.expand_nested_zips()?;
73 }
74 Ok(archive)
75 }
76
77 fn expand_nested_zips(&mut self) -> ImportResult<()> {
78 let mut expanded = Vec::new();
79 for entry in self.entries.drain(..).collect::<Vec<_>>() {
80 if entry.meta.path.to_lowercase().ends_with(".zip") {
81 let base_path = strip_extension(&entry.meta.path).to_string();
82 let bytes = if let Some(bytes) = entry.bytes {
83 bytes
84 } else {
85 self.source.read_entry(entry.meta.index)?
86 };
87 let nested = read_zip_bytes_entries(&bytes, &base_path)?;
88 expanded.extend(nested);
89 } else {
90 expanded.push(entry);
91 }
92 }
93 self.entries = expanded;
94 Ok(())
95 }
96
97 fn read_entry(&self, entry: &IndexedEntry) -> ImportResult<Vec<u8>> {
98 if let Some(bytes) = &entry.bytes {
99 return Ok(bytes.clone());
100 }
101 self.source.read_entry(entry.meta.index)
102 }
103}
104
105fn read_zip_bytes_entries(bytes: &[u8], base_path: &str) -> ImportResult<Vec<IndexedEntry>> {
106 use std::io::{Cursor, Read};
107
108 let mut archive = zip::ZipArchive::new(Cursor::new(bytes))?;
109 let mut entries = Vec::new();
110 for index in 0..archive.len() {
111 let mut file = archive.by_index(index)?;
112 if file.is_dir() {
113 continue;
114 }
115 let Some(path) = file
116 .enclosed_name()
117 .map(|path| normalize_import_path(&path.to_string_lossy()))
118 else {
119 continue;
120 };
121 if path.is_empty() || crate::source::is_system_path(&path) {
122 continue;
123 }
124 let mut bytes = Vec::with_capacity(file.size() as usize);
125 file.read_to_end(&mut bytes)?;
126 let path = if base_path.is_empty() {
127 path
128 } else {
129 format!("{base_path}/{path}")
130 };
131 entries.push(IndexedEntry {
132 meta: ArchiveEntryMeta {
133 index,
134 path,
135 compressed_size: file.compressed_size(),
136 uncompressed_size: file.size(),
137 },
138 bytes: Some(bytes),
139 });
140 }
141 Ok(entries)
142}
143
144#[derive(Debug, Clone)]
145struct ImportAssetRef {
146 entry_index: usize,
147 blob_id: String,
148}
149
150#[derive(Debug, Clone)]
151struct ImportDocRef {
152 entry_index: usize,
153 doc_id: String,
154 title: Option<String>,
155 icon: Option<ImportedIconData>,
156}
157
158fn read_archive(
159 source: ImportSource,
160 expand_nested_zips: bool,
161 options: &ImportOptions,
162) -> ImportResult<IndexedArchive> {
163 if options.cancel {
164 return Err(ImportError::Cancelled);
165 }
166 match source {
167 ImportSource::FilePath(path) => IndexedArchive::read(ZipPathSource::new(path), expand_nested_zips),
168 ImportSource::DirectoryPath(path) => IndexedArchive::read(DirectoryPathSource::new(path), expand_nested_zips),
169 }
170}
171
172fn empty_batch(total: usize) -> ImportBatch {
173 ImportBatch {
174 docs: Vec::new(),
175 blobs: Vec::new(),
176 folders: Vec::new(),
177 tags: Vec::new(),
178 icons: Vec::new(),
179 warnings: Vec::new(),
180 progress: ImportProgress { completed: 0, total },
181 entry_id: None,
182 is_workspace_file: false,
183 done: false,
184 }
185}
186
187fn hash_bytes(bytes: &[u8]) -> String {
188 let mut hasher = Sha256::new();
189 hasher.update(bytes);
190 hasher
191 .finalize()
192 .iter()
193 .map(|byte| format!("{byte:02x}"))
194 .collect::<String>()
195}
196
197fn entry_from_indexed(entry: &IndexedEntry, bytes: Vec<u8>) -> VfsEntry {
198 VfsEntry {
199 path: entry.meta.path.clone(),
200 bytes,
201 }
202}
203
204fn asset_from_indexed(
205 archive: &IndexedArchive,
206 entry: &IndexedEntry,
207 asset: &ImportAssetRef,
208) -> ImportResult<ImportedAsset> {
209 Ok(ImportedAsset {
210 blob_id: asset.blob_id.clone(),
211 source_path: entry.meta.path.clone(),
212 file_name: file_name(&entry.meta.path).to_string(),
213 mime: mime_from_path(&entry.meta.path).to_string(),
214 bytes: archive.read_entry(entry)?,
215 })
216}
217
218fn add_asset_batch(
219 archive: &IndexedArchive,
220 limits: &ImportBatchLimits,
221 emitted_assets: &mut BTreeSet<usize>,
222 batch: &mut ImportBatch,
223 assets: &[ImportAssetRef],
224) -> ImportResult<()> {
225 let mut bytes_in_batch = 0u64;
226 let max_blobs = limits.max_blobs.max(1);
227 for (asset_index, asset) in assets.iter().enumerate() {
228 if emitted_assets.contains(&asset_index) || batch.blobs.len() >= max_blobs {
229 continue;
230 }
231 let entry = &archive.entries[asset.entry_index];
232 if entry.meta.uncompressed_size > limits.max_blob_bytes {
233 batch.warnings.push(ImportWarning {
234 code: "skipped_asset".to_string(),
235 source_path: Some(entry.meta.path.clone()),
236 message: format!(
237 "Skipped {}: asset is larger than the batch import support",
238 entry.meta.path
239 ),
240 });
241 emitted_assets.insert(asset_index);
242 continue;
243 }
244 if !batch.blobs.is_empty() && bytes_in_batch + entry.meta.uncompressed_size > limits.max_blob_bytes {
245 continue;
246 }
247 let imported = asset_from_indexed(archive, entry, asset)?;
248 bytes_in_batch += entry.meta.uncompressed_size;
249 batch.blobs.push(imported);
250 emitted_assets.insert(asset_index);
251 }
252 Ok(())
253}
254
255fn has_emitted_assets(emitted_assets: &BTreeSet<usize>, assets: &[ImportAssetRef]) -> bool {
256 emitted_assets.len() >= assets.len()
257}
258
259fn folders_for_doc_path(path: &str, doc_id: &str, icon: Option<ImportedIconData>) -> Vec<FolderHierarchyDelta> {
260 folder_hierarchy_deltas(folder_parts(path), Some(doc_id), icon, |folder| folder.to_string())
261}
262
263fn folder_hierarchy_deltas<I, F>(
264 parts: I,
265 doc_id: Option<&str>,
266 icon: Option<ImportedIconData>,
267 mut name_for_part: F,
268) -> Vec<FolderHierarchyDelta>
269where
270 I: IntoIterator,
271 I::Item: AsRef<str>,
272 F: FnMut(&str) -> String,
273{
274 let mut folders = Vec::new();
275 let mut current_path = String::new();
276 for part in parts {
277 let folder = part.as_ref();
278 let parent_path = (!current_path.is_empty()).then(|| current_path.clone());
279 current_path = if current_path.is_empty() {
280 folder.to_string()
281 } else {
282 format!("{current_path}/{folder}")
283 };
284 folders.push(FolderHierarchyDelta {
285 path: current_path.clone(),
286 name: name_for_part(folder),
287 parent_path,
288 page_id: None,
289 icon: None,
290 });
291 }
292 if let Some(doc_id) = doc_id
293 && !current_path.is_empty()
294 {
295 folders.push(FolderHierarchyDelta {
296 path: format!("{current_path}/__doc__{doc_id}"),
297 name: format!("__doc__{doc_id}"),
298 parent_path: Some(current_path),
299 page_id: Some(doc_id.to_string()),
300 icon,
301 });
302 }
303 folders
304}
305
306fn doc_snapshot(
307 doc_id: &str,
308 title: &str,
309 markdown: &str,
310 current_path: &str,
311 page_ids_by_path: &BTreeMap<String, String>,
312 blob_ids_by_path: &BTreeMap<String, String>,
313) -> ImportResult<serde_json::Value> {
314 doc_snapshot_inner(
315 doc_id,
316 title,
317 markdown,
318 current_path,
319 page_ids_by_path,
320 blob_ids_by_path,
321 None,
322 )
323}
324
325fn doc_snapshot_with_id_hints(
326 doc_id: &str,
327 title: &str,
328 markdown: &str,
329 current_path: &str,
330 page_ids_by_path: &BTreeMap<String, String>,
331 blob_ids_by_path: &BTreeMap<String, String>,
332 id_hints: (&str, &str),
333) -> ImportResult<serde_json::Value> {
334 doc_snapshot_inner(
335 doc_id,
336 title,
337 markdown,
338 current_path,
339 page_ids_by_path,
340 blob_ids_by_path,
341 Some(id_hints),
342 )
343}
344
345fn doc_snapshot_inner(
346 doc_id: &str,
347 title: &str,
348 markdown: &str,
349 current_path: &str,
350 page_ids_by_path: &BTreeMap<String, String>,
351 blob_ids_by_path: &BTreeMap<String, String>,
352 id_hints: Option<(&str, &str)>,
353) -> ImportResult<serde_json::Value> {
354 let rewritten = rewrite_markdown_link_destinations(markdown, current_path, page_ids_by_path);
355 let rewritten = rewrite_bare_page_references(&rewritten, current_path, page_ids_by_path);
356 let rewritten = rewrite_asset_references(&rewritten, current_path, blob_ids_by_path);
357 let mut snapshot = match id_hints {
358 Some((namespace, token)) => {
359 affine_doc_loader::build_doc_snapshot_with_id_hints(title, &rewritten, doc_id, namespace, token)?
360 }
361 None => affine_doc_loader::build_doc_snapshot(title, &rewritten, doc_id)?,
362 };
363 rewrite_page_references(&mut snapshot, current_path, page_ids_by_path);
364 Ok(snapshot)
365}
366
367fn push_skipped_doc_warning(batch: &mut ImportBatch, source_path: &str, error: affine_doc_loader::ParseError) {
368 batch.warnings.push(ImportWarning {
369 code: "skipped_doc".to_string(),
370 source_path: Some(source_path.to_string()),
371 message: format!("Skipped {source_path}: {error}"),
372 });
373}
374
375fn folder_parts(path: &str) -> Vec<String> {
376 let mut parts = normalize_import_path(path)
377 .split('/')
378 .map(ToString::to_string)
379 .collect::<Vec<_>>();
380 parts.pop();
381 if parts.len() > 1 {
382 parts.remove(0);
383 }
384 parts
385}