1use std::sync::Arc;
9
10use async_trait::async_trait;
11use bytes::Bytes;
12use futures::StreamExt;
13use lance_core::utils::tracing::{
14 AUDIT_MODE_CREATE, AUDIT_MODE_DELETE, AUDIT_TYPE_MANIFEST, TRACE_FILE_AUDIT,
15};
16use lance_core::{Error, Result};
17use lance_io::object_store::ObjectStore;
18use log::warn;
19use object_store::ObjectMeta;
20use object_store::ObjectStoreExt;
21use object_store::{Error as ObjectStoreError, ObjectStore as OSObjectStore, path::Path};
22use tracing::info;
23
24use super::{
25 MANIFEST_EXTENSION, ManifestLocation, ManifestNamingScheme, current_manifest_path,
26 default_resolve_version, make_staging_manifest_path, write_version_hint,
27};
28use crate::format::{IndexMetadata, Manifest, Transaction};
29use crate::io::commit::{CommitError, CommitHandler};
30
31#[async_trait]
44pub trait ExternalManifestStore: std::fmt::Debug + Send + Sync {
45 async fn get(&self, base_uri: &str, version: u64) -> Result<String>;
47
48 async fn get_manifest_location(
49 &self,
50 base_uri: &str,
51 version: u64,
52 ) -> Result<ManifestLocation> {
53 let path = self.get(base_uri, version).await?;
54 let path = Path::parse(&path).map_err(|e| Error::invalid_input(e.to_string()))?;
55 let naming_scheme = detect_naming_scheme_from_path(&path)?;
56 Ok(ManifestLocation {
57 version,
58 path,
59 size: None,
60 naming_scheme,
61 e_tag: None,
62 })
63 }
64
65 async fn get_latest_version(&self, base_uri: &str) -> Result<Option<(u64, String)>>;
69
70 async fn get_latest_manifest_location(
76 &self,
77 base_uri: &str,
78 ) -> Result<Option<ManifestLocation>> {
79 self.get_latest_version(base_uri).await.and_then(|res| {
80 res.map(|(version, uri)| {
81 let path = Path::parse(&uri).map_err(|e| Error::invalid_input(e.to_string()))?;
82 let naming_scheme = detect_naming_scheme_from_path(&path)?;
83 Ok(ManifestLocation {
84 version,
85 path,
86 size: None,
87 naming_scheme,
88 e_tag: None,
89 })
90 })
91 .transpose()
92 })
93 }
94
95 #[allow(clippy::too_many_arguments)]
104 async fn put(
105 &self,
106 base_path: &Path,
107 version: u64,
108 staging_path: &Path,
109 size: u64,
110 e_tag: Option<String>,
111 object_store: &dyn OSObjectStore,
112 naming_scheme: ManifestNamingScheme,
113 ) -> Result<ManifestLocation> {
114 self.put_if_not_exists(
118 base_path.as_ref(),
119 version,
120 staging_path.as_ref(),
121 size,
122 e_tag.clone(),
123 )
124 .await?;
125
126 let final_path = naming_scheme.manifest_path(base_path, version);
128 let copied = match copy_size_aware(object_store, staging_path, &final_path, size).await {
129 Ok(_) => true,
130 Err(ObjectStoreError::NotFound { .. }) => false,
131 Err(e) => return Err(e.into()),
132 };
133 if copied {
134 info!(target: TRACE_FILE_AUDIT, mode=AUDIT_MODE_CREATE, r#type=AUDIT_TYPE_MANIFEST, path = final_path.as_ref());
135 }
136
137 let e_tag = if copied && size < 5 * 1024 * 1024 {
139 e_tag
140 } else {
141 let meta = object_store.head(&final_path).await?;
142 meta.e_tag
143 };
144
145 let location = ManifestLocation {
146 version,
147 path: final_path.clone(),
148 size: Some(size),
149 naming_scheme,
150 e_tag: e_tag.clone(),
151 };
152
153 if !copied {
154 return Ok(location);
155 }
156
157 self.put_if_exists(
159 base_path.as_ref(),
160 version,
161 final_path.as_ref(),
162 size,
163 e_tag,
164 )
165 .await?;
166
167 match object_store.delete(staging_path).await {
169 Ok(_) => {}
170 Err(ObjectStoreError::NotFound { .. }) => {}
171 Err(e) => return Err(e.into()),
172 }
173 info!(target: TRACE_FILE_AUDIT, mode=AUDIT_MODE_DELETE, r#type=AUDIT_TYPE_MANIFEST, path = staging_path.as_ref());
174
175 Ok(location)
176 }
177
178 async fn put_if_not_exists(
180 &self,
181 base_uri: &str,
182 version: u64,
183 path: &str,
184 size: u64,
185 e_tag: Option<String>,
186 ) -> Result<()>;
187
188 async fn put_if_exists(
190 &self,
191 base_uri: &str,
192 version: u64,
193 path: &str,
194 size: u64,
195 e_tag: Option<String>,
196 ) -> Result<()>;
197
198 async fn delete(&self, _base_uri: &str) -> Result<()> {
200 Ok(())
201 }
202}
203
204pub(crate) fn detect_naming_scheme_from_path(path: &Path) -> Result<ManifestNamingScheme> {
205 path.filename()
206 .and_then(|name| {
207 ManifestNamingScheme::detect_scheme(name)
208 .or_else(|| Some(ManifestNamingScheme::detect_scheme_staging(name)))
209 })
210 .ok_or_else(|| {
211 Error::corrupt_file(
212 path.clone(),
213 "Path does not follow known manifest naming convention.",
214 )
215 })
216}
217
218const MAX_SERVER_SIDE_COPY_BYTES: u64 = 5 * 1024 * 1024 * 1024;
227
228const COPY_REWRITE_PART_SIZE: usize = 100 * 1024 * 1024;
234
235async fn copy_size_aware(
256 store: &dyn OSObjectStore,
257 from: &Path,
258 to: &Path,
259 size: u64,
260) -> std::result::Result<(), ObjectStoreError> {
261 if size < MAX_SERVER_SIDE_COPY_BYTES {
262 store.copy(from, to).await
263 } else {
264 copy_via_read_rewrite(store, from, to).await
265 }
266}
267
268async fn copy_via_read_rewrite(
275 store: &dyn OSObjectStore,
276 from: &Path,
277 to: &Path,
278) -> std::result::Result<(), ObjectStoreError> {
279 let mut stream = store.get(from).await?.into_stream();
281
282 let mut upload = store.put_multipart(to).await?;
292 let mut part_buf: Vec<u8> = Vec::with_capacity(COPY_REWRITE_PART_SIZE);
293
294 while let Some(chunk) = stream.next().await {
295 let chunk = match chunk {
296 Ok(b) => b,
297 Err(e) => {
298 let _ = upload.abort().await;
299 return Err(e);
300 }
301 };
302 let mut offset = 0;
308 while offset < chunk.len() {
309 let want = COPY_REWRITE_PART_SIZE - part_buf.len();
310 let take = want.min(chunk.len() - offset);
311 part_buf.extend_from_slice(&chunk[offset..offset + take]);
312 offset += take;
313
314 if part_buf.len() >= COPY_REWRITE_PART_SIZE {
315 let payload =
316 std::mem::replace(&mut part_buf, Vec::with_capacity(COPY_REWRITE_PART_SIZE));
317 if let Err(e) = upload.put_part(Bytes::from(payload).into()).await {
318 let _ = upload.abort().await;
319 return Err(e);
320 }
321 }
322 }
323 }
324
325 if !part_buf.is_empty()
328 && let Err(e) = upload.put_part(Bytes::from(part_buf).into()).await
329 {
330 let _ = upload.abort().await;
331 return Err(e);
332 }
333
334 if let Err(e) = upload.complete().await {
335 let _ = upload.abort().await;
336 return Err(e);
337 }
338 Ok(())
339}
340
341#[derive(Debug)]
345pub struct ExternalManifestCommitHandler {
346 pub external_manifest_store: Arc<dyn ExternalManifestStore>,
347}
348
349impl ExternalManifestCommitHandler {
350 async fn verify_finalized_manifest_location(
351 &self,
352 base_path: &Path,
353 location: ManifestLocation,
354 object_store: &dyn OSObjectStore,
355 ) -> std::result::Result<ManifestLocation, Error> {
356 match object_store.head(&location.path).await {
357 Ok(ObjectMeta { size, e_tag, .. }) => {
358 let ManifestLocation {
359 version,
360 path,
361 size: expected_size,
362 naming_scheme,
363 e_tag: expected_e_tag,
364 } = location;
365
366 let size = match expected_size {
367 Some(expected_size) if expected_size != size => {
368 return Err(Error::corrupt_file(
369 path,
370 format!(
371 "Manifest size mismatch for version {}: external store expected {}, object store returned {}",
372 version, expected_size, size
373 ),
374 ));
375 }
376 Some(expected_size) => Some(expected_size),
377 None => Some(size),
378 };
379
380 let e_tag = match expected_e_tag {
381 Some(expected_e_tag) => {
382 if e_tag.as_ref() != Some(&expected_e_tag) {
383 return Err(Error::corrupt_file(
384 path,
385 format!(
386 "Manifest e_tag mismatch for version {}: external store expected {:?}, object store returned {:?}",
387 version, expected_e_tag, e_tag
388 ),
389 ));
390 }
391 Some(expected_e_tag)
392 }
393 None => e_tag,
394 };
395
396 Ok(ManifestLocation {
397 version,
398 path,
399 size,
400 naming_scheme,
401 e_tag,
402 })
403 }
404 Err(ObjectStoreError::NotFound { .. }) => {
405 default_resolve_version(base_path, location.version, object_store).await
408 }
409 Err(e) => Err(e.into()),
410 }
411 }
412
413 #[allow(clippy::too_many_arguments)]
423 async fn finalize_manifest(
424 &self,
425 base_path: &Path,
426 staging_manifest_path: &Path,
427 version: u64,
428 size: u64,
429 e_tag: Option<String>,
430 store: &dyn OSObjectStore,
431 naming_scheme: ManifestNamingScheme,
432 ) -> std::result::Result<ManifestLocation, Error> {
433 let final_manifest_path = naming_scheme.manifest_path(base_path, version);
435
436 let copied =
437 match copy_size_aware(store, staging_manifest_path, &final_manifest_path, size).await {
438 Ok(_) => true,
439 Err(ObjectStoreError::NotFound { .. }) => false, Err(e) => return Err(e.into()),
441 };
442 if copied {
443 info!(target: TRACE_FILE_AUDIT, mode=AUDIT_MODE_CREATE, r#type=AUDIT_TYPE_MANIFEST, path = final_manifest_path.as_ref());
444 }
445
446 let e_tag = if copied && size < 5 * 1024 * 1024 {
453 e_tag
454 } else {
455 let meta = store.head(&final_manifest_path).await?;
456 meta.e_tag
457 };
458
459 let location = ManifestLocation {
460 version,
461 path: final_manifest_path,
462 size: Some(size),
463 naming_scheme,
464 e_tag,
465 };
466
467 if !copied {
468 return Ok(location);
469 }
470
471 self.external_manifest_store
473 .put_if_exists(
474 base_path.as_ref(),
475 version,
476 location.path.as_ref(),
477 size,
478 location.e_tag.clone(),
479 )
480 .await?;
481
482 match store.delete(staging_manifest_path).await {
484 Ok(_) => {}
485 Err(ObjectStoreError::NotFound { .. }) => {}
486 Err(e) => return Err(e.into()),
487 }
488 info!(target: TRACE_FILE_AUDIT, mode=AUDIT_MODE_DELETE, r#type=AUDIT_TYPE_MANIFEST, path = staging_manifest_path.as_ref());
489
490 Ok(location)
491 }
492}
493
494#[async_trait]
495impl CommitHandler for ExternalManifestCommitHandler {
496 async fn resolve_latest_location(
497 &self,
498 base_path: &Path,
499 object_store: &ObjectStore,
500 ) -> std::result::Result<ManifestLocation, Error> {
501 let location = self
502 .external_manifest_store
503 .get_latest_manifest_location(base_path.as_ref())
504 .await?;
505
506 match location {
507 Some(location) => {
508 if location.path.extension() == Some(MANIFEST_EXTENSION) {
509 return self
510 .verify_finalized_manifest_location(
511 base_path,
512 location,
513 object_store.inner.as_ref(),
514 )
515 .await;
516 }
517
518 let ManifestLocation {
519 version,
520 path,
521 size,
522 naming_scheme,
523 e_tag,
524 } = location;
525
526 let (size, e_tag) = if let Some(size) = size {
527 (size, e_tag)
528 } else {
529 match object_store.inner.head(&path).await {
530 Ok(meta) => (meta.size, meta.e_tag),
531 Err(ObjectStoreError::NotFound { .. }) => {
532 let new_location = self
534 .external_manifest_store
535 .get_manifest_location(base_path.as_ref(), version)
536 .await?;
537 return Ok(new_location);
538 }
539 Err(e) => return Err(e.into()),
540 }
541 };
542
543 let final_location = self
544 .finalize_manifest(
545 base_path,
546 &path,
547 version,
548 size,
549 e_tag.clone(),
550 &object_store.inner,
551 naming_scheme,
552 )
553 .await?;
554
555 Ok(final_location)
556 }
557 None => current_manifest_path(object_store, base_path).await,
560 }
561 }
562
563 async fn resolve_version_location(
564 &self,
565 base_path: &Path,
566 version: u64,
567 object_store: &dyn OSObjectStore,
568 ) -> std::result::Result<ManifestLocation, Error> {
569 let location_res = self
570 .external_manifest_store
571 .get_manifest_location(base_path.as_ref(), version)
572 .await;
573
574 let location = match location_res {
575 Ok(p) => p,
576 Err(Error::NotFound { .. }) => {
578 let path = default_resolve_version(base_path, version, object_store)
579 .await
580 .map_err(|_| Error::not_found(format!("{}@{}", base_path, version)))?
581 .path;
582 match object_store.head(&path).await {
583 Ok(ObjectMeta { size, e_tag, .. }) => {
584 let res = self
585 .external_manifest_store
586 .put_if_not_exists(
587 base_path.as_ref(),
588 version,
589 path.as_ref(),
590 size,
591 e_tag.clone(),
592 )
593 .await;
594 if let Err(e) = res {
595 warn!(
596 "could not update external manifest store during load, with error: {}",
597 e
598 );
599 }
600 let naming_scheme =
601 ManifestNamingScheme::detect_scheme_staging(path.filename().unwrap());
602 return Ok(ManifestLocation {
603 version,
604 path,
605 size: Some(size),
606 naming_scheme,
607 e_tag,
608 });
609 }
610 Err(ObjectStoreError::NotFound { .. }) => {
611 return Err(Error::not_found(path.to_string()));
612 }
613 Err(e) => return Err(e.into()),
614 }
615 }
616 Err(e) => return Err(e),
617 };
618
619 if location.path.extension() == Some(MANIFEST_EXTENSION) {
620 return self
621 .verify_finalized_manifest_location(base_path, location, object_store)
622 .await;
623 }
624
625 let naming_scheme =
626 ManifestNamingScheme::detect_scheme_staging(location.path.filename().unwrap());
627
628 let (size, e_tag) = if let Some(size) = location.size {
629 (size, location.e_tag.clone())
630 } else {
631 let meta = object_store.head(&location.path).await?;
632 (meta.size as u64, meta.e_tag)
633 };
634
635 self.finalize_manifest(
636 base_path,
637 &location.path,
638 version,
639 size,
640 e_tag,
641 object_store,
642 naming_scheme,
643 )
644 .await
645 }
646
647 async fn version_exists(
648 &self,
649 base_path: &Path,
650 version: u64,
651 object_store: &dyn OSObjectStore,
652 naming_scheme: ManifestNamingScheme,
653 ) -> Result<bool> {
654 match self
655 .external_manifest_store
656 .get_manifest_location(base_path.as_ref(), version)
657 .await
658 {
659 Ok(_) => Ok(true),
660 Err(Error::NotFound { .. }) => {
661 let path = naming_scheme.manifest_path(base_path, version);
662 match object_store.head(&path).await {
663 Ok(_) => Ok(true),
664 Err(ObjectStoreError::NotFound { .. }) => Ok(false),
665 Err(e) => Err(e.into()),
666 }
667 }
668 Err(e) => Err(e),
669 }
670 }
671
672 async fn commit(
673 &self,
674 manifest: &mut Manifest,
675 indices: Option<Vec<IndexMetadata>>,
676 base_path: &Path,
677 object_store: &ObjectStore,
678 manifest_writer: super::ManifestWriter,
679 naming_scheme: ManifestNamingScheme,
680 transaction: Option<Transaction>,
681 ) -> std::result::Result<ManifestLocation, CommitError> {
682 let path = naming_scheme.manifest_path(base_path, manifest.version);
687 let staging_path = make_staging_manifest_path(&path)?;
688 let write_res =
689 manifest_writer(object_store, manifest, indices, &staging_path, transaction).await?;
690
691 let result = self
693 .external_manifest_store
694 .put(
695 base_path,
696 manifest.version,
697 &staging_path,
698 write_res.size as u64,
699 write_res.e_tag,
700 &object_store.inner,
701 naming_scheme,
702 )
703 .await;
704
705 match result {
706 Ok(location) => {
707 write_version_hint(object_store, base_path, manifest.version).await;
708 Ok(location)
709 }
710 Err(_) => {
711 match object_store.inner.delete(&staging_path).await {
713 Ok(_) => {}
714 Err(ObjectStoreError::NotFound { .. }) => {}
715 Err(e) => return Err(CommitError::OtherError(e.into())),
716 }
717 info!(target: TRACE_FILE_AUDIT, mode=AUDIT_MODE_DELETE, r#type=AUDIT_TYPE_MANIFEST, path = staging_path.as_ref());
718 Err(CommitError::CommitConflict {})
719 }
720 }
721 }
722
723 async fn delete(&self, base_path: &Path) -> Result<()> {
724 self.external_manifest_store
725 .delete(base_path.as_ref())
726 .await
727 }
728}