1use std::io;
26use std::pin::Pin;
27use std::sync::Arc;
28use std::sync::atomic::AtomicBool;
29use std::{fmt::Debug, fs::DirEntry};
30
31use super::manifest::write_manifest;
32use futures::Stream;
33use futures::future::Either;
34use futures::{
35 StreamExt, TryStreamExt,
36 future::{self, BoxFuture},
37 stream::BoxStream,
38};
39use lance_file::format::{MAGIC, MAJOR_VERSION, MINOR_VERSION};
40use lance_io::object_writer::{ObjectWriter, WriteResult, get_etag};
41use log::warn;
42use object_store::ObjectStoreExt as OSObjectStoreExt;
43use object_store::PutOptions;
44use object_store::{Error as ObjectStoreError, ObjectStore as OSObjectStore, path::Path};
45use tracing::info;
46use url::Url;
47
48#[cfg(feature = "dynamodb")]
49pub mod dynamodb;
50pub mod external_manifest;
51
52use lance_core::{Error, Result};
53use lance_io::object_store::{ObjectStore, ObjectStoreExt, ObjectStoreParams};
54use lance_io::traits::{WriteExt, Writer};
55
56use crate::format::{IndexMetadata, Manifest, Transaction, is_detached_version};
57use lance_core::utils::tracing::{AUDIT_MODE_CREATE, AUDIT_TYPE_MANIFEST, TRACE_FILE_AUDIT};
58#[cfg(feature = "dynamodb")]
59use {
60 self::external_manifest::{ExternalManifestCommitHandler, ExternalManifestStore},
61 aws_credential_types::provider::ProvideCredentials,
62 aws_credential_types::provider::error::CredentialsError,
63 lance_io::object_store::{StorageOptions, providers::aws::build_aws_credential},
64 object_store::aws::AmazonS3ConfigKey,
65 object_store::aws::AwsCredentialProvider,
66 std::borrow::Cow,
67 std::time::{Duration, SystemTime},
68};
69
70pub const VERSIONS_DIR: &str = "_versions";
71const MANIFEST_EXTENSION: &str = "manifest";
72const DETACHED_VERSION_PREFIX: &str = "d";
73const VERSION_HINT_FILE: &str = "latest_version_hint.json";
80
81#[derive(Clone, Copy, Debug, PartialEq, Eq)]
83pub enum ManifestNamingScheme {
84 V1,
86 V2,
90}
91
92impl ManifestNamingScheme {
93 pub fn manifest_path(&self, base: &Path, version: u64) -> Path {
94 if is_detached_version(version) {
95 base.clone().join(VERSIONS_DIR).join(format!(
100 "{DETACHED_VERSION_PREFIX}{version}.{MANIFEST_EXTENSION}"
101 ))
102 } else {
103 let directory = base.clone().join(VERSIONS_DIR);
104 match self {
105 Self::V1 => directory.join(format!("{version}.{MANIFEST_EXTENSION}")),
106 Self::V2 => {
107 let inverted_version = u64::MAX - version;
108 directory.join(format!("{inverted_version:020}.{MANIFEST_EXTENSION}"))
109 }
110 }
111 }
112 }
113
114 pub fn parse_version(&self, filename: &str) -> Option<u64> {
115 let file_number = filename
116 .split_once('.')
117 .and_then(|(version_str, _)| version_str.parse::<u64>().ok());
119 match self {
120 Self::V1 => file_number,
121 Self::V2 => file_number.map(|v| u64::MAX - v),
122 }
123 }
124
125 pub fn parse_detached_version(filename: &str) -> Option<u64> {
129 if !filename.starts_with(DETACHED_VERSION_PREFIX) {
130 return None;
131 }
132 let without_prefix = &filename[DETACHED_VERSION_PREFIX.len()..];
133 without_prefix
134 .split_once('.')
135 .and_then(|(version_str, _)| version_str.parse::<u64>().ok())
136 }
137
138 pub fn detect_scheme(filename: &str) -> Option<Self> {
139 if filename.starts_with(DETACHED_VERSION_PREFIX) {
140 return Some(Self::V2);
142 }
143 if filename.ends_with(MANIFEST_EXTENSION) {
144 const V2_LEN: usize = 20 + 1 + MANIFEST_EXTENSION.len();
145 if filename.len() == V2_LEN {
146 Some(Self::V2)
147 } else {
148 Some(Self::V1)
149 }
150 } else {
151 None
152 }
153 }
154
155 pub fn detect_scheme_staging(filename: &str) -> Self {
156 if filename.chars().nth(20) == Some('.') {
159 Self::V2
160 } else {
161 Self::V1
162 }
163 }
164}
165
166pub async fn migrate_scheme_to_v2(object_store: &ObjectStore, dataset_base: &Path) -> Result<()> {
176 object_store
177 .inner
178 .list(Some(&dataset_base.clone().join(VERSIONS_DIR)))
179 .try_filter(|res| {
180 let res = if let Some(filename) = res.location.filename() {
181 ManifestNamingScheme::detect_scheme(filename) == Some(ManifestNamingScheme::V1)
182 } else {
183 false
184 };
185 future::ready(res)
186 })
187 .try_for_each_concurrent(object_store.io_parallelism(), |meta| async move {
188 let filename = meta.location.filename().unwrap();
189 let version = ManifestNamingScheme::V1.parse_version(filename).unwrap();
190 let path = ManifestNamingScheme::V2.manifest_path(dataset_base, version);
191 object_store.inner.rename(&meta.location, &path).await?;
192 Ok(())
193 })
194 .await?;
195
196 Ok(())
197}
198
199pub type ManifestWriter = for<'a> fn(
203 object_store: &'a ObjectStore,
204 manifest: &'a mut Manifest,
205 indices: Option<Vec<IndexMetadata>>,
206 path: &'a Path,
207 transaction: Option<Transaction>,
208) -> BoxFuture<'a, Result<WriteResult>>;
209
210pub fn write_manifest_file_to_path<'a>(
214 object_store: &'a ObjectStore,
215 manifest: &'a mut Manifest,
216 indices: Option<Vec<IndexMetadata>>,
217 path: &'a Path,
218 transaction: Option<Transaction>,
219) -> BoxFuture<'a, Result<WriteResult>> {
220 Box::pin(async move {
221 let mut object_writer = ObjectWriter::new(object_store, path).await?;
222 let pos = write_manifest(&mut object_writer, manifest, indices, transaction).await?;
223 object_writer
224 .write_magics(pos, MAJOR_VERSION, MINOR_VERSION, MAGIC)
225 .await?;
226 let res = Writer::shutdown(&mut object_writer).await?;
227 info!(target: TRACE_FILE_AUDIT, mode=AUDIT_MODE_CREATE, r#type=AUDIT_TYPE_MANIFEST, path = path.to_string());
228 Ok(res)
229 })
230}
231
232#[derive(Debug, Clone)]
233pub struct ManifestLocation {
234 pub version: u64,
236 pub path: Path,
238 pub size: Option<u64>,
240 pub naming_scheme: ManifestNamingScheme,
242 pub e_tag: Option<String>,
260}
261
262impl TryFrom<object_store::ObjectMeta> for ManifestLocation {
263 type Error = Error;
264
265 fn try_from(meta: object_store::ObjectMeta) -> Result<Self> {
266 let filename = meta.location.filename().ok_or_else(|| {
267 Error::internal("ObjectMeta location does not have a filename".to_string())
268 })?;
269 let scheme = ManifestNamingScheme::detect_scheme(filename)
270 .ok_or_else(|| Error::internal(format!("Invalid manifest filename: '{}'", filename)))?;
271 let version = scheme
272 .parse_version(filename)
273 .ok_or_else(|| Error::internal(format!("Invalid manifest filename: '{}'", filename)))?;
274 Ok(Self {
275 version,
276 path: meta.location,
277 size: Some(meta.size),
278 naming_scheme: scheme,
279 e_tag: meta.e_tag,
280 })
281 }
282}
283
284async fn current_manifest_path(
294 object_store: &ObjectStore,
295 base: &Path,
296) -> Result<ManifestLocation> {
297 if object_store.has_direct_local_paths() {
298 if let Ok(Some(location)) = current_manifest_local(base) {
299 return Ok(location);
300 }
301 } else if uses_version_hint(object_store)
302 && let Some(location) = read_version_hint_and_probe(object_store, base).await
303 {
304 return Ok(location);
305 }
306
307 resolve_version_from_listing(object_store, base).await
308}
309
310#[derive(serde::Serialize, serde::Deserialize)]
312struct VersionHint {
313 version: u64,
314}
315
316const VERSION_HINT_ENV: &str = "LANCE_USE_VERSION_HINT";
321
322fn version_hint_globally_enabled() -> bool {
323 static ENABLED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
324 *ENABLED.get_or_init(|| match std::env::var(VERSION_HINT_ENV) {
325 Ok(v) => !matches!(
326 v.trim().to_ascii_lowercase().as_str(),
327 "0" | "false" | "off"
328 ),
329 Err(_) => true,
330 })
331}
332
333pub fn uses_version_hint(object_store: &ObjectStore) -> bool {
342 version_hint_globally_enabled() && !object_store.list_is_lexically_ordered
343}
344
345fn version_hint_path(base: &Path) -> Path {
347 base.clone().join(VERSIONS_DIR).join(VERSION_HINT_FILE)
348}
349
350pub async fn write_version_hint(object_store: &ObjectStore, base: &Path, version: u64) {
358 if is_detached_version(version) || !uses_version_hint(object_store) {
359 return;
360 }
361 let hint_path = version_hint_path(base);
362 let content = serde_json::to_vec(&VersionHint { version }).expect("serialize version hint");
363 if let Err(e) = object_store.put(&hint_path, content.as_slice()).await {
364 warn!("Failed to write version hint file for version {version}: {e}");
365 }
366}
367
368async fn read_version_from_hint(object_store: &ObjectStore, base: &Path) -> Option<u64> {
371 let bytes = object_store
372 .inner
373 .get(&version_hint_path(base))
374 .await
375 .ok()?
376 .bytes()
377 .await
378 .ok()?;
379 Some(serde_json::from_slice::<VersionHint>(&bytes).ok()?.version)
380}
381
382async fn read_version_hint_and_probe(
387 object_store: &ObjectStore,
388 base: &Path,
389) -> Option<ManifestLocation> {
390 let hint_version = read_version_from_hint(object_store, base).await?;
391 let (version, scheme, mut probed) = probe_versions_upward(object_store, base, hint_version)
392 .await
393 .ok()
394 .flatten()?;
395 let (_, meta) = probed.pop()?;
397 Some(ManifestLocation {
398 version,
399 path: scheme.manifest_path(base, version),
400 size: Some(meta.size),
401 naming_scheme: scheme,
402 e_tag: meta.e_tag,
403 })
404}
405
406const MAX_HINT_PROBE_GAP: u64 = 1000;
410
411async fn probe_versions_upward(
428 object_store: &ObjectStore,
429 base: &Path,
430 from_version: u64,
431) -> Result<
432 Option<(
433 u64,
434 ManifestNamingScheme,
435 Vec<(u64, object_store::ObjectMeta)>,
436 )>,
437> {
438 let mut scheme = ManifestNamingScheme::V2;
440 let meta = match object_store
441 .inner
442 .head(&scheme.manifest_path(base, from_version))
443 .await
444 {
445 Ok(meta) => meta,
446 Err(ObjectStoreError::NotFound { .. }) => {
447 scheme = ManifestNamingScheme::V1;
448 match object_store
449 .inner
450 .head(&scheme.manifest_path(base, from_version))
451 .await
452 {
453 Ok(meta) => meta,
454 Err(ObjectStoreError::NotFound { .. }) => return Ok(None),
455 Err(e) => return Err(e.into()),
456 }
457 }
458 Err(e) => return Err(e.into()),
459 };
460
461 let mut probed = vec![(from_version, meta)];
462 let mut version = from_version;
463 loop {
464 let next = version + 1;
465 match object_store
466 .inner
467 .head(&scheme.manifest_path(base, next))
468 .await
469 {
470 Ok(meta) => {
471 probed.push((next, meta));
472 version = next;
473 }
474 Err(ObjectStoreError::NotFound { .. }) => break,
476 Err(e) => return Err(e.into()),
479 }
480 }
481 Ok(Some((version, scheme, probed)))
482}
483
484async fn list_manifests_since_version_with_hint(
491 object_store: &ObjectStore,
492 base: &Path,
493 since_version: u64,
494) -> Option<Vec<ManifestLocation>> {
495 let hint_version = read_version_from_hint(object_store, base).await?;
496
497 if hint_version.saturating_sub(since_version) > MAX_HINT_PROBE_GAP {
500 return None;
501 }
502
503 let probe_from = if hint_version > since_version {
506 hint_version
507 } else {
508 since_version + 1
509 };
510
511 let (scheme, probed) = match probe_versions_upward(object_store, base, probe_from).await {
512 Ok(Some((_true_latest, scheme, probed))) => (scheme, probed),
513 Ok(None) if hint_version > since_version => return None,
517 Ok(None) => return Some(Vec::new()),
518 Err(_) => return None,
520 };
521
522 let mut locations: Vec<ManifestLocation> = probed
523 .into_iter()
524 .filter(|(v, _)| *v > since_version)
525 .map(|(version, meta)| ManifestLocation {
526 version,
527 path: scheme.manifest_path(base, version),
528 size: Some(meta.size),
529 naming_scheme: scheme,
530 e_tag: meta.e_tag,
531 })
532 .collect();
533
534 if hint_version > since_version + 1 {
539 let gap_locations: Vec<ManifestLocation> =
540 futures::stream::iter((since_version + 1)..hint_version)
541 .map(|version| async move {
542 object_store
543 .inner
544 .head(&scheme.manifest_path(base, version))
545 .await
546 .map(|meta| ManifestLocation {
547 version,
548 path: scheme.manifest_path(base, version),
549 size: Some(meta.size),
550 naming_scheme: scheme,
551 e_tag: meta.e_tag,
552 })
553 })
554 .buffer_unordered(object_store.io_parallelism())
555 .try_collect()
556 .await
557 .ok()?;
558 locations.extend(gap_locations);
559 }
560
561 locations.sort_by_key(|loc| std::cmp::Reverse(loc.version));
562 Some(locations)
563}
564
565async fn resolve_version_from_listing(
567 object_store: &ObjectStore,
568 base: &Path,
569) -> Result<ManifestLocation> {
570 let manifest_files = object_store.list(Some(base.clone().join(VERSIONS_DIR)));
571
572 let mut valid_manifests = manifest_files.try_filter_map(|res| {
573 let filename = res.location.filename().unwrap();
574 if let Some(scheme) = ManifestNamingScheme::detect_scheme(filename) {
575 if scheme.parse_version(filename).is_some() {
577 future::ready(Ok(Some((scheme, res))))
578 } else {
579 future::ready(Ok(None))
580 }
581 } else {
582 future::ready(Ok(None))
583 }
584 });
585
586 let first = valid_manifests.next().await.transpose()?;
587 match (first, object_store.list_is_lexically_ordered) {
588 (Some((scheme @ ManifestNamingScheme::V2, meta)), true) => {
591 let version = scheme
592 .parse_version(meta.location.filename().unwrap())
593 .unwrap();
594
595 for (scheme, meta) in valid_manifests.take(999).try_collect::<Vec<_>>().await? {
599 if scheme != ManifestNamingScheme::V2 {
600 warn!(
601 "Found V1 Manifest in a V2 directory. Use `migrate_manifest_paths_v2` \
602 to migrate the directory."
603 );
604 break;
605 }
606 let next_version = scheme
607 .parse_version(meta.location.filename().unwrap())
608 .unwrap();
609 if next_version >= version {
610 warn!(
611 "List operation was expected to be lexically ordered, but was not. This \
612 could mean a corrupt read. Please make a bug report on the lance-format/lance \
613 GitHub repository."
614 );
615 break;
616 }
617 }
618
619 Ok(ManifestLocation {
620 version,
621 path: meta.location,
622 size: Some(meta.size),
623 naming_scheme: scheme,
624 e_tag: meta.e_tag,
625 })
626 }
627 (Some((first_scheme, meta)), _) => {
630 let mut current_version = first_scheme
631 .parse_version(meta.location.filename().unwrap())
632 .unwrap();
633 let mut current_meta = meta;
634 let scheme = first_scheme;
635
636 while let Some((entry_scheme, meta)) = valid_manifests.next().await.transpose()? {
637 if entry_scheme != scheme {
638 return Err(Error::internal(format!(
639 "Found multiple manifest naming schemes in the same directory: {:?} and {:?}. \
640 Use `migrate_manifest_paths_v2` to migrate the directory.",
641 scheme, entry_scheme
642 )));
643 }
644 let version = entry_scheme
645 .parse_version(meta.location.filename().unwrap())
646 .unwrap();
647 if version > current_version {
648 current_version = version;
649 current_meta = meta;
650 }
651 }
652 Ok(ManifestLocation {
653 version: current_version,
654 path: current_meta.location,
655 size: Some(current_meta.size),
656 naming_scheme: scheme,
657 e_tag: current_meta.e_tag,
658 })
659 }
660 (None, _) => Err(Error::not_found(
661 base.clone().join(VERSIONS_DIR).to_string(),
662 )),
663 }
664}
665
666fn current_manifest_local(base: &Path) -> std::io::Result<Option<ManifestLocation>> {
670 let path = lance_io::local::to_local_path(&base.clone().join(VERSIONS_DIR));
671 let entries = std::fs::read_dir(path)?;
672
673 let mut latest_entry: Option<(u64, DirEntry, ManifestNamingScheme)> = None;
674
675 let mut scheme: Option<ManifestNamingScheme> = None;
676
677 for entry in entries {
678 let entry = entry?;
679 let filename_raw = entry.file_name();
680 let filename = filename_raw.to_string_lossy();
681
682 let Some(entry_scheme) = ManifestNamingScheme::detect_scheme(&filename) else {
683 continue;
686 };
687
688 if let Some(scheme) = scheme {
689 if scheme != entry_scheme {
690 return Err(io::Error::new(
691 io::ErrorKind::InvalidData,
692 format!(
693 "Found multiple manifest naming schemes in the same directory: {:?} and {:?}",
694 scheme, entry_scheme
695 ),
696 ));
697 }
698 } else {
699 scheme = Some(entry_scheme);
700 }
701
702 let Some(version) = entry_scheme.parse_version(&filename) else {
703 continue;
704 };
705
706 if let Some((latest_version, _, _)) = &latest_entry {
707 if version > *latest_version {
708 latest_entry = Some((version, entry, entry_scheme));
709 }
710 } else {
711 latest_entry = Some((version, entry, entry_scheme));
712 }
713 }
714
715 if let Some((version, entry, naming_scheme)) = latest_entry {
716 let metadata = entry.metadata()?;
717 Ok(Some(ManifestLocation {
718 version,
719 path: naming_scheme.manifest_path(base, version),
720 size: Some(metadata.len()),
721 naming_scheme,
722 e_tag: Some(get_etag(&metadata)),
723 }))
724 } else {
725 Ok(None)
726 }
727}
728
729fn list_manifests<'a>(
730 base_path: &Path,
731 object_store: &'a dyn OSObjectStore,
732) -> impl Stream<Item = Result<ManifestLocation>> + 'a {
733 object_store
734 .read_dir_all(&base_path.clone().join(VERSIONS_DIR), None)
735 .filter_map(|obj_meta| {
736 futures::future::ready(
737 obj_meta
738 .map(|m| ManifestLocation::try_from(m).ok())
739 .transpose(),
740 )
741 })
742 .boxed()
743}
744
745fn detached_manifest_location_from_meta(
747 meta: object_store::ObjectMeta,
748) -> Option<ManifestLocation> {
749 let filename = meta.location.filename()?;
750 let version = ManifestNamingScheme::parse_detached_version(filename)?;
751 Some(ManifestLocation {
752 version,
753 path: meta.location,
754 size: Some(meta.size),
755 naming_scheme: ManifestNamingScheme::V2,
756 e_tag: meta.e_tag,
757 })
758}
759
760pub fn list_detached_manifests<'a>(
762 base_path: &Path,
763 object_store: &'a dyn OSObjectStore,
764) -> impl Stream<Item = Result<ManifestLocation>> + 'a {
765 object_store
766 .read_dir_all(&base_path.clone().join(VERSIONS_DIR), None)
767 .filter_map(|obj_meta| {
768 futures::future::ready(
769 obj_meta
770 .map(detached_manifest_location_from_meta)
771 .transpose(),
772 )
773 })
774 .boxed()
775}
776
777fn make_staging_manifest_path(base: &Path) -> Result<Path> {
778 let id = uuid::Uuid::new_v4().to_string();
779 Path::parse(format!("{base}-{id}")).map_err(|e| Error::io_source(Box::new(e)))
780}
781
782#[cfg(feature = "dynamodb")]
783const DDB_URL_QUERY_KEY: &str = "ddbTableName";
784
785#[async_trait::async_trait]
794#[allow(clippy::too_many_arguments)]
795pub trait CommitHandler: Debug + Send + Sync {
796 fn is_version_not_found_definitive(&self) -> bool {
804 false
805 }
806
807 fn propagate_commit_error_after_success(&self) -> bool {
814 true
815 }
816
817 async fn resolve_latest_location(
818 &self,
819 base_path: &Path,
820 object_store: &ObjectStore,
821 ) -> Result<ManifestLocation> {
822 Ok(current_manifest_path(object_store, base_path).await?)
823 }
824
825 async fn resolve_version_location(
826 &self,
827 base_path: &Path,
828 version: u64,
829 object_store: &dyn OSObjectStore,
830 ) -> Result<ManifestLocation> {
831 default_resolve_version(base_path, version, object_store).await
832 }
833
834 async fn version_exists(
840 &self,
841 base_path: &Path,
842 version: u64,
843 object_store: &dyn OSObjectStore,
844 naming_scheme: ManifestNamingScheme,
845 ) -> Result<bool> {
846 let path = naming_scheme.manifest_path(base_path, version);
847 match object_store.head(&path).await {
848 Ok(_) => Ok(true),
849 Err(ObjectStoreError::NotFound { .. }) => Ok(false),
850 Err(e) => Err(e.into()),
851 }
852 }
853
854 fn list_detached_manifest_locations<'a>(
858 &self,
859 base_path: &Path,
860 object_store: &'a ObjectStore,
861 ) -> BoxStream<'a, Result<ManifestLocation>> {
862 list_detached_manifests(base_path, &object_store.inner).boxed()
863 }
864
865 fn list_manifest_locations<'a>(
872 &self,
873 base_path: &Path,
874 object_store: &'a ObjectStore,
875 sorted_descending: bool,
876 ) -> BoxStream<'a, Result<ManifestLocation>> {
877 let underlying_stream = list_manifests(base_path, &object_store.inner);
878
879 if !sorted_descending {
880 return underlying_stream.boxed();
881 }
882
883 async fn sort_stream(
884 input_stream: impl futures::Stream<Item = Result<ManifestLocation>> + Unpin,
885 ) -> Result<impl Stream<Item = Result<ManifestLocation>> + Unpin> {
886 let mut locations = input_stream.try_collect::<Vec<_>>().await?;
887 locations.sort_by_key(|m| std::cmp::Reverse(m.version));
888 Ok(futures::stream::iter(locations.into_iter().map(Ok)))
889 }
890
891 if object_store.list_is_lexically_ordered {
894 let mut peekable = underlying_stream.peekable();
896
897 futures::stream::once(async move {
898 let naming_scheme = match Pin::new(&mut peekable).peek().await {
899 Some(Ok(m)) => m.naming_scheme,
900 Some(Err(_)) => ManifestNamingScheme::V2,
903 None => ManifestNamingScheme::V2,
904 };
905
906 if naming_scheme == ManifestNamingScheme::V2 {
907 Ok(Either::Left(peekable))
909 } else {
910 sort_stream(peekable).await.map(Either::Right)
911 }
912 })
913 .try_flatten()
914 .boxed()
915 } else {
916 futures::stream::once(sort_stream(underlying_stream))
921 .try_flatten()
922 .boxed()
923 }
924 }
925
926 fn list_manifest_locations_since<'a>(
934 &self,
935 base_path: &Path,
936 object_store: &'a ObjectStore,
937 since_version: u64,
938 ) -> BoxStream<'a, Result<ManifestLocation>> {
939 if !uses_version_hint(object_store) {
940 return self
941 .list_manifest_locations(base_path, object_store, true)
942 .try_take_while(move |loc| future::ready(Ok(loc.version > since_version)))
943 .boxed();
944 }
945
946 let base_path = base_path.clone();
947 futures::stream::once(async move {
948 let locations = match list_manifests_since_version_with_hint(
949 object_store,
950 &base_path,
951 since_version,
952 )
953 .await
954 {
955 Some(locations) => locations,
956 None => {
957 let mut locations = list_manifests(&base_path, &object_store.inner)
958 .try_collect::<Vec<_>>()
959 .await?;
960 locations.retain(|loc| loc.version > since_version);
961 locations.sort_by_key(|loc| std::cmp::Reverse(loc.version));
962 locations
963 }
964 };
965 Ok::<_, Error>(futures::stream::iter(locations.into_iter().map(Ok)))
966 })
967 .try_flatten()
968 .boxed()
969 }
970
971 async fn commit(
976 &self,
977 manifest: &mut Manifest,
978 indices: Option<Vec<IndexMetadata>>,
979 base_path: &Path,
980 object_store: &ObjectStore,
981 manifest_writer: ManifestWriter,
982 naming_scheme: ManifestNamingScheme,
983 transaction: Option<Transaction>,
984 ) -> std::result::Result<ManifestLocation, CommitError>;
985
986 async fn delete(&self, _base_path: &Path) -> Result<()> {
988 Ok(())
989 }
990}
991
992async fn default_resolve_version(
993 base_path: &Path,
994 version: u64,
995 object_store: &dyn OSObjectStore,
996) -> Result<ManifestLocation> {
997 if is_detached_version(version) {
998 return Ok(ManifestLocation {
999 version,
1000 naming_scheme: ManifestNamingScheme::V2,
1003 path: ManifestNamingScheme::V2.manifest_path(base_path, version),
1005 size: None,
1006 e_tag: None,
1007 });
1008 }
1009
1010 let scheme = ManifestNamingScheme::V2;
1012 let path = scheme.manifest_path(base_path, version);
1013 match object_store.head(&path).await {
1014 Ok(meta) => Ok(ManifestLocation {
1015 version,
1016 path,
1017 size: Some(meta.size),
1018 naming_scheme: scheme,
1019 e_tag: meta.e_tag,
1020 }),
1021 Err(ObjectStoreError::NotFound { .. }) => {
1022 let scheme = ManifestNamingScheme::V1;
1024 Ok(ManifestLocation {
1025 version,
1026 path: scheme.manifest_path(base_path, version),
1027 size: None,
1028 naming_scheme: scheme,
1029 e_tag: None,
1030 })
1031 }
1032 Err(e) => Err(e.into()),
1033 }
1034}
1035#[cfg(feature = "dynamodb")]
1037#[derive(Debug)]
1038struct OSObjectStoreToAwsCredAdaptor(AwsCredentialProvider);
1039
1040#[cfg(feature = "dynamodb")]
1041impl ProvideCredentials for OSObjectStoreToAwsCredAdaptor {
1042 fn provide_credentials<'a>(
1043 &'a self,
1044 ) -> aws_credential_types::provider::future::ProvideCredentials<'a>
1045 where
1046 Self: 'a,
1047 {
1048 aws_credential_types::provider::future::ProvideCredentials::new(async {
1049 let creds = self
1050 .0
1051 .get_credential()
1052 .await
1053 .map_err(|e| CredentialsError::provider_error(Box::new(e)))?;
1054 Ok(aws_credential_types::Credentials::new(
1055 &creds.key_id,
1056 &creds.secret_key,
1057 creds.token.clone(),
1058 Some(
1059 SystemTime::now()
1060 .checked_add(Duration::from_secs(
1061 60 * 10, ))
1063 .expect("overflow"),
1064 ),
1065 "",
1066 ))
1067 })
1068 }
1069}
1070
1071#[cfg(feature = "dynamodb")]
1072async fn build_dynamodb_external_store(
1073 table_name: &str,
1074 creds: AwsCredentialProvider,
1075 region: &str,
1076 endpoint: Option<String>,
1077 app_name: &str,
1078) -> Result<Arc<dyn ExternalManifestStore>> {
1079 use super::commit::dynamodb::DynamoDBExternalManifestStore;
1080 use aws_sdk_dynamodb::{
1081 Client,
1082 config::{IdentityCache, Region, retry::RetryConfig},
1083 };
1084
1085 let mut dynamodb_config = aws_sdk_dynamodb::config::Builder::new()
1086 .behavior_version_latest()
1087 .region(Some(Region::new(region.to_string())))
1088 .credentials_provider(OSObjectStoreToAwsCredAdaptor(creds))
1089 .identity_cache(IdentityCache::no_cache())
1091 .retry_config(RetryConfig::standard().with_max_attempts(5));
1094
1095 if let Some(endpoint) = endpoint {
1096 dynamodb_config = dynamodb_config.endpoint_url(endpoint);
1097 }
1098 let client = Client::from_conf(dynamodb_config.build());
1099
1100 DynamoDBExternalManifestStore::new_external_store(client.into(), table_name, app_name).await
1101}
1102
1103pub async fn commit_handler_from_url(
1104 url_or_path: &str,
1105 #[allow(unused_variables)] options: &Option<ObjectStoreParams>,
1107) -> Result<Arc<dyn CommitHandler>> {
1108 let local_handler: Arc<dyn CommitHandler> = if cfg!(windows) {
1109 Arc::new(RenameCommitHandler)
1110 } else {
1111 Arc::new(ConditionalPutCommitHandler)
1112 };
1113
1114 let url = match Url::parse(url_or_path) {
1115 Ok(url) if url.scheme().len() == 1 && cfg!(windows) => {
1116 return Ok(local_handler);
1118 }
1119 Ok(url) => url,
1120 Err(_) => {
1121 return Ok(local_handler);
1122 }
1123 };
1124
1125 match url.scheme() {
1126 "file" | "file-object-store" => Ok(local_handler),
1127 "s3" | "gs" | "az" | "abfss" | "memory" | "oss" | "tos" | "shared-memory" | "goosefs" => {
1128 Ok(Arc::new(ConditionalPutCommitHandler))
1129 }
1130 "cos" => Ok(Arc::new(TencentCosCommitHandler)),
1131 #[cfg(not(feature = "dynamodb"))]
1132 "s3+ddb" => Err(Error::invalid_input_source(
1133 "`s3+ddb://` scheme requires `dynamodb` feature to be enabled".into(),
1134 )),
1135 #[cfg(feature = "dynamodb")]
1136 "s3+ddb" => {
1137 if url.query_pairs().count() != 1 {
1138 return Err(Error::invalid_input_source(
1139 "`s3+ddb://` scheme and expects exactly one query `ddbTableName`".into(),
1140 ));
1141 }
1142 let table_name = match url.query_pairs().next() {
1143 Some((Cow::Borrowed(key), Cow::Borrowed(table_name)))
1144 if key == DDB_URL_QUERY_KEY =>
1145 {
1146 if table_name.is_empty() {
1147 return Err(Error::invalid_input_source(
1148 "`s3+ddb://` scheme requires non empty dynamodb table name".into(),
1149 ));
1150 }
1151 table_name
1152 }
1153 _ => {
1154 return Err(Error::invalid_input_source(
1155 "`s3+ddb://` scheme and expects exactly one query `ddbTableName`".into(),
1156 ));
1157 }
1158 };
1159 let options = options.clone().unwrap_or_default();
1160 let storage_options_raw =
1161 StorageOptions(options.storage_options().cloned().unwrap_or_default());
1162 let dynamo_endpoint = get_dynamodb_endpoint(&storage_options_raw);
1163 let storage_options = storage_options_raw.as_s3_options();
1164
1165 let region = storage_options.get(&AmazonS3ConfigKey::Region).cloned();
1166
1167 let accessor = options.get_accessor();
1169
1170 let provider_scheme = storage_options_raw.aws_provider_scheme()?;
1171
1172 let (aws_creds, region) = build_aws_credential(
1173 options.s3_credentials_refresh_offset,
1174 options.aws_credentials.clone(),
1175 Some(&storage_options),
1176 region,
1177 accessor,
1178 provider_scheme,
1179 )
1180 .await?;
1181
1182 Ok(Arc::new(ExternalManifestCommitHandler {
1183 external_manifest_store: build_dynamodb_external_store(
1184 table_name,
1185 aws_creds.clone(),
1186 ®ion,
1187 dynamo_endpoint,
1188 "lancedb",
1189 )
1190 .await?,
1191 }))
1192 }
1193 _ => Ok(Arc::new(UnsafeCommitHandler)),
1194 }
1195}
1196
1197#[cfg(feature = "dynamodb")]
1198fn get_dynamodb_endpoint(storage_options: &StorageOptions) -> Option<String> {
1199 if let Some(endpoint) = storage_options.0.get("dynamodb_endpoint") {
1200 Some(endpoint.clone())
1201 } else {
1202 std::env::var("DYNAMODB_ENDPOINT").ok()
1203 }
1204}
1205
1206#[derive(Debug)]
1208pub enum CommitError {
1209 CommitConflict,
1211 OtherError(Error),
1213}
1214
1215impl From<Error> for CommitError {
1216 fn from(e: Error) -> Self {
1217 Self::OtherError(e)
1218 }
1219}
1220
1221impl From<CommitError> for Error {
1222 fn from(e: CommitError) -> Self {
1223 match e {
1224 CommitError::CommitConflict => Self::internal("Commit conflict".to_string()),
1225 CommitError::OtherError(e) => e,
1226 }
1227 }
1228}
1229
1230static WARNED_ON_UNSAFE_COMMIT: AtomicBool = AtomicBool::new(false);
1232
1233pub struct UnsafeCommitHandler;
1237
1238#[async_trait::async_trait]
1239#[allow(clippy::too_many_arguments)]
1240impl CommitHandler for UnsafeCommitHandler {
1241 fn is_version_not_found_definitive(&self) -> bool {
1242 true
1243 }
1244
1245 fn propagate_commit_error_after_success(&self) -> bool {
1246 false
1247 }
1248
1249 async fn commit(
1250 &self,
1251 manifest: &mut Manifest,
1252 indices: Option<Vec<IndexMetadata>>,
1253 base_path: &Path,
1254 object_store: &ObjectStore,
1255 manifest_writer: ManifestWriter,
1256 naming_scheme: ManifestNamingScheme,
1257 transaction: Option<Transaction>,
1258 ) -> std::result::Result<ManifestLocation, CommitError> {
1259 if !WARNED_ON_UNSAFE_COMMIT.load(std::sync::atomic::Ordering::Relaxed) {
1261 WARNED_ON_UNSAFE_COMMIT.store(true, std::sync::atomic::Ordering::Relaxed);
1262 log::warn!(
1263 "Using unsafe commit handler. Concurrent writes may result in data loss. \
1264 Consider providing a commit handler that prevents conflicting writes."
1265 );
1266 }
1267
1268 let version_path = naming_scheme.manifest_path(base_path, manifest.version);
1269 let res =
1270 manifest_writer(object_store, manifest, indices, &version_path, transaction).await?;
1271
1272 write_version_hint(object_store, base_path, manifest.version).await;
1273
1274 Ok(ManifestLocation {
1275 version: manifest.version,
1276 size: Some(res.size as u64),
1277 naming_scheme,
1278 path: version_path,
1279 e_tag: res.e_tag,
1280 })
1281 }
1282}
1283
1284impl Debug for UnsafeCommitHandler {
1285 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1286 f.debug_struct("UnsafeCommitHandler").finish()
1287 }
1288}
1289
1290#[async_trait::async_trait]
1292pub trait CommitLock: Debug {
1293 type Lease: CommitLease;
1294
1295 async fn lock(&self, version: u64) -> std::result::Result<Self::Lease, CommitError>;
1308}
1309
1310#[async_trait::async_trait]
1311pub trait CommitLease: Send + Sync {
1312 async fn release(&self, success: bool) -> std::result::Result<(), CommitError>;
1318}
1319
1320struct LeaseGuard<L: CommitLease + 'static> {
1329 lease: Option<L>,
1330}
1331
1332impl<L: CommitLease + 'static> LeaseGuard<L> {
1333 fn new(lease: L) -> Self {
1334 Self { lease: Some(lease) }
1335 }
1336
1337 async fn release(mut self, success: bool) -> std::result::Result<(), CommitError> {
1339 let result = {
1344 let lease = self
1345 .lease
1346 .as_ref()
1347 .expect("LeaseGuard released more than once");
1348 lease.release(success).await
1349 };
1350 self.lease = None;
1351 result
1352 }
1353}
1354
1355impl<L: CommitLease + 'static> Drop for LeaseGuard<L> {
1356 fn drop(&mut self) {
1357 if let Some(lease) = self.lease.take() {
1358 if let Ok(handle) = tokio::runtime::Handle::try_current() {
1363 handle.spawn(async move {
1364 let _ = lease.release(false).await;
1365 });
1366 }
1367 }
1368 }
1369}
1370
1371#[async_trait::async_trait]
1372impl<T: CommitLock + Send + Sync> CommitHandler for T
1373where
1374 T::Lease: 'static,
1375{
1376 fn is_version_not_found_definitive(&self) -> bool {
1377 true
1378 }
1379
1380 async fn commit(
1381 &self,
1382 manifest: &mut Manifest,
1383 indices: Option<Vec<IndexMetadata>>,
1384 base_path: &Path,
1385 object_store: &ObjectStore,
1386 manifest_writer: ManifestWriter,
1387 naming_scheme: ManifestNamingScheme,
1388 transaction: Option<Transaction>,
1389 ) -> std::result::Result<ManifestLocation, CommitError> {
1390 let path = naming_scheme.manifest_path(base_path, manifest.version);
1391 let lease = LeaseGuard::new(self.lock(manifest.version).await?);
1396
1397 match object_store.inner.head(&path).await {
1399 Ok(_) => {
1400 lease.release(false).await?;
1403
1404 return Err(CommitError::CommitConflict);
1405 }
1406 Err(ObjectStoreError::NotFound { .. }) => {}
1407 Err(e) => {
1408 lease.release(false).await?;
1411
1412 return Err(CommitError::OtherError(e.into()));
1413 }
1414 }
1415 let res = manifest_writer(object_store, manifest, indices, &path, transaction).await;
1416
1417 lease.release(res.is_ok()).await?;
1419
1420 let res = res?;
1421
1422 write_version_hint(object_store, base_path, manifest.version).await;
1423
1424 Ok(ManifestLocation {
1425 version: manifest.version,
1426 size: Some(res.size as u64),
1427 naming_scheme,
1428 path,
1429 e_tag: res.e_tag,
1430 })
1431 }
1432}
1433
1434#[async_trait::async_trait]
1435impl<T: CommitLock + Send + Sync> CommitHandler for Arc<T>
1436where
1437 T::Lease: 'static,
1438{
1439 fn is_version_not_found_definitive(&self) -> bool {
1440 self.as_ref().is_version_not_found_definitive()
1441 }
1442
1443 fn propagate_commit_error_after_success(&self) -> bool {
1444 self.as_ref().propagate_commit_error_after_success()
1445 }
1446
1447 async fn commit(
1448 &self,
1449 manifest: &mut Manifest,
1450 indices: Option<Vec<IndexMetadata>>,
1451 base_path: &Path,
1452 object_store: &ObjectStore,
1453 manifest_writer: ManifestWriter,
1454 naming_scheme: ManifestNamingScheme,
1455 transaction: Option<Transaction>,
1456 ) -> std::result::Result<ManifestLocation, CommitError> {
1457 self.as_ref()
1458 .commit(
1459 manifest,
1460 indices,
1461 base_path,
1462 object_store,
1463 manifest_writer,
1464 naming_scheme,
1465 transaction,
1466 )
1467 .await
1468 }
1469}
1470
1471pub struct RenameCommitHandler;
1475
1476#[async_trait::async_trait]
1477impl CommitHandler for RenameCommitHandler {
1478 fn is_version_not_found_definitive(&self) -> bool {
1479 true
1480 }
1481
1482 fn propagate_commit_error_after_success(&self) -> bool {
1483 false
1484 }
1485
1486 async fn commit(
1487 &self,
1488 manifest: &mut Manifest,
1489 indices: Option<Vec<IndexMetadata>>,
1490 base_path: &Path,
1491 object_store: &ObjectStore,
1492 manifest_writer: ManifestWriter,
1493 naming_scheme: ManifestNamingScheme,
1494 transaction: Option<Transaction>,
1495 ) -> std::result::Result<ManifestLocation, CommitError> {
1496 let path = naming_scheme.manifest_path(base_path, manifest.version);
1500 let tmp_path = make_staging_manifest_path(&path)?;
1501
1502 let res = manifest_writer(object_store, manifest, indices, &tmp_path, transaction).await?;
1503
1504 match object_store
1505 .inner
1506 .rename_if_not_exists(&tmp_path, &path)
1507 .await
1508 {
1509 Ok(_) => {
1510 write_version_hint(object_store, base_path, manifest.version).await;
1512 Ok(ManifestLocation {
1513 version: manifest.version,
1514 path,
1515 size: Some(res.size as u64),
1516 naming_scheme,
1517 e_tag: None, })
1519 }
1520 Err(ObjectStoreError::AlreadyExists { .. }) => {
1521 let _ = object_store.delete(&tmp_path).await;
1524
1525 return Err(CommitError::CommitConflict);
1526 }
1527 Err(e) => {
1528 return Err(CommitError::OtherError(e.into()));
1530 }
1531 }
1532 }
1533}
1534
1535impl Debug for RenameCommitHandler {
1536 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1537 f.debug_struct("RenameCommitHandler").finish()
1538 }
1539}
1540
1541pub struct ConditionalPutCommitHandler;
1542
1543#[async_trait::async_trait]
1544impl CommitHandler for ConditionalPutCommitHandler {
1545 fn is_version_not_found_definitive(&self) -> bool {
1546 true
1547 }
1548
1549 fn propagate_commit_error_after_success(&self) -> bool {
1550 false
1551 }
1552
1553 async fn commit(
1554 &self,
1555 manifest: &mut Manifest,
1556 indices: Option<Vec<IndexMetadata>>,
1557 base_path: &Path,
1558 object_store: &ObjectStore,
1559 manifest_writer: ManifestWriter,
1560 naming_scheme: ManifestNamingScheme,
1561 transaction: Option<Transaction>,
1562 ) -> std::result::Result<ManifestLocation, CommitError> {
1563 let path = naming_scheme.manifest_path(base_path, manifest.version);
1564
1565 let memory_store = ObjectStore::memory();
1566 let dummy_path = "dummy";
1567 manifest_writer(
1568 &memory_store,
1569 manifest,
1570 indices,
1571 &dummy_path.into(),
1572 transaction,
1573 )
1574 .await?;
1575 let dummy_data = memory_store.read_one_all(&dummy_path.into()).await?;
1576 let size = dummy_data.len() as u64;
1577 let res = object_store
1578 .inner
1579 .put_opts(
1580 &path,
1581 dummy_data.into(),
1582 PutOptions {
1583 mode: object_store::PutMode::Create,
1584 ..Default::default()
1585 },
1586 )
1587 .await
1588 .map_err(|err| match err {
1589 ObjectStoreError::AlreadyExists { .. } | ObjectStoreError::Precondition { .. } => {
1590 CommitError::CommitConflict
1591 }
1592 _ => CommitError::OtherError(err.into()),
1593 })?;
1594
1595 write_version_hint(object_store, base_path, manifest.version).await;
1596
1597 Ok(ManifestLocation {
1598 version: manifest.version,
1599 path,
1600 size: Some(size),
1601 naming_scheme,
1602 e_tag: res.e_tag,
1603 })
1604 }
1605}
1606
1607impl Debug for ConditionalPutCommitHandler {
1608 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1609 f.debug_struct("ConditionalPutCommitHandler").finish()
1610 }
1611}
1612
1613struct TencentCosCommitHandler;
1621
1622#[async_trait::async_trait]
1623impl CommitHandler for TencentCosCommitHandler {
1624 fn is_version_not_found_definitive(&self) -> bool {
1625 true
1626 }
1627
1628 async fn commit(
1629 &self,
1630 _manifest: &mut Manifest,
1631 _indices: Option<Vec<IndexMetadata>>,
1632 _base_path: &Path,
1633 _object_store: &ObjectStore,
1634 _manifest_writer: ManifestWriter,
1635 _naming_scheme: ManifestNamingScheme,
1636 _transaction: Option<Transaction>,
1637 ) -> std::result::Result<ManifestLocation, CommitError> {
1638 Err(CommitError::OtherError(Error::not_supported(
1639 "Default writes to Tencent COS are disabled because COS does not reliably enforce \
1640 put-if-not-exists after bucket versioning has ever been enabled. Provide a \
1641 distributed commit_lock in Python or a custom CommitHandler in Rust.",
1642 )))
1643 }
1644}
1645
1646impl Debug for TencentCosCommitHandler {
1647 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1648 f.debug_struct("TencentCosCommitHandler").finish()
1649 }
1650}
1651
1652#[derive(Debug, Clone)]
1653pub struct CommitConfig {
1654 pub num_retries: u32,
1655 pub skip_auto_cleanup: bool,
1656 }
1658
1659impl Default for CommitConfig {
1660 fn default() -> Self {
1661 Self {
1662 num_retries: 20,
1663 skip_auto_cleanup: false,
1664 }
1665 }
1666}
1667
1668#[cfg(test)]
1669mod tests {
1670 use std::sync::atomic::AtomicUsize;
1671
1672 use lance_core::utils::tempfile::TempObjDir;
1673
1674 use super::*;
1675
1676 #[test]
1677 fn test_manifest_naming_scheme() {
1678 let v1 = ManifestNamingScheme::V1;
1679 let v2 = ManifestNamingScheme::V2;
1680
1681 assert_eq!(
1682 v1.manifest_path(&Path::from("base"), 0),
1683 Path::from("base/_versions/0.manifest")
1684 );
1685 assert_eq!(
1686 v1.manifest_path(&Path::from("base"), 42),
1687 Path::from("base/_versions/42.manifest")
1688 );
1689
1690 assert_eq!(
1691 v2.manifest_path(&Path::from("base"), 0),
1692 Path::from("base/_versions/18446744073709551615.manifest")
1693 );
1694 assert_eq!(
1695 v2.manifest_path(&Path::from("base"), 42),
1696 Path::from("base/_versions/18446744073709551573.manifest")
1697 );
1698
1699 assert_eq!(v1.parse_version("0.manifest"), Some(0));
1700 assert_eq!(v1.parse_version("42.manifest"), Some(42));
1701 assert_eq!(
1702 v1.parse_version("42.manifest-cee4fbbb-eb19-4ea3-8ca7-54f5ec33dedc"),
1703 Some(42)
1704 );
1705
1706 assert_eq!(v2.parse_version("18446744073709551615.manifest"), Some(0));
1707 assert_eq!(v2.parse_version("18446744073709551573.manifest"), Some(42));
1708 assert_eq!(
1709 v2.parse_version("18446744073709551573.manifest-cee4fbbb-eb19-4ea3-8ca7-54f5ec33dedc"),
1710 Some(42)
1711 );
1712
1713 assert_eq!(ManifestNamingScheme::detect_scheme("0.manifest"), Some(v1));
1714 assert_eq!(
1715 ManifestNamingScheme::detect_scheme("18446744073709551615.manifest"),
1716 Some(v2)
1717 );
1718 assert_eq!(ManifestNamingScheme::detect_scheme("something else"), None);
1719 }
1720
1721 #[tokio::test]
1722 async fn test_manifest_naming_migration() {
1723 let object_store = ObjectStore::memory();
1724 let base = Path::from("base");
1725 let versions_dir = base.clone().join(VERSIONS_DIR);
1726
1727 let original_files = vec![
1729 versions_dir.clone().join("irrelevant"),
1730 ManifestNamingScheme::V1.manifest_path(&base, 0),
1731 ManifestNamingScheme::V2.manifest_path(&base, 1),
1732 ];
1733 for path in original_files {
1734 object_store.put(&path, b"".as_slice()).await.unwrap();
1735 }
1736
1737 migrate_scheme_to_v2(&object_store, &base).await.unwrap();
1738
1739 let expected_files = vec![
1740 ManifestNamingScheme::V2.manifest_path(&base, 1),
1741 ManifestNamingScheme::V2.manifest_path(&base, 0),
1742 versions_dir.clone().join("irrelevant"),
1743 ];
1744 let actual_files = object_store
1745 .inner
1746 .list(Some(&versions_dir))
1747 .map_ok(|res| res.location)
1748 .try_collect::<Vec<_>>()
1749 .await
1750 .unwrap();
1751 assert_eq!(actual_files, expected_files);
1752 }
1753
1754 #[tokio::test]
1755 #[rstest::rstest]
1756 async fn test_list_manifests_sorted(
1757 #[values(true, false)] lexical_list_store: bool,
1758 #[values(ManifestNamingScheme::V1, ManifestNamingScheme::V2)]
1759 naming_scheme: ManifestNamingScheme,
1760 ) {
1761 let tempdir;
1762 let (object_store, base) = if lexical_list_store {
1763 (Box::new(ObjectStore::memory()), Path::from("base"))
1764 } else {
1765 tempdir = TempObjDir::default();
1766 let path = tempdir.clone().join("base");
1767 let store = Box::new(ObjectStore::local());
1768 assert!(!store.list_is_lexically_ordered);
1769 (store, path)
1770 };
1771
1772 let mut expected_paths = Vec::new();
1774 for i in (0..12).rev() {
1775 let path = naming_scheme.manifest_path(&base, i);
1776 object_store.put(&path, b"".as_slice()).await.unwrap();
1777 expected_paths.push(path);
1778 }
1779
1780 let actual_versions = ConditionalPutCommitHandler
1781 .list_manifest_locations(&base, &object_store, true)
1782 .map_ok(|location| location.path)
1783 .try_collect::<Vec<_>>()
1784 .await
1785 .unwrap();
1786
1787 assert_eq!(actual_versions, expected_paths);
1788 }
1789
1790 #[tokio::test]
1791 #[rstest::rstest]
1792 async fn test_current_manifest_path(
1793 #[values(true, false)] lexical_list_store: bool,
1794 #[values(ManifestNamingScheme::V1, ManifestNamingScheme::V2)]
1795 naming_scheme: ManifestNamingScheme,
1796 ) {
1797 let mut object_store = ObjectStore::memory();
1800 object_store.list_is_lexically_ordered = lexical_list_store;
1801 let object_store = Box::new(object_store);
1802 let base = Path::from("base");
1803
1804 for version in [5, 2, 11, 0, 8, 3, 10, 1, 7, 4, 9, 6] {
1806 let path = naming_scheme.manifest_path(&base, version);
1807 object_store.put(&path, b"".as_slice()).await.unwrap();
1808 }
1809
1810 let location = current_manifest_path(&object_store, &base).await.unwrap();
1811
1812 assert_eq!(location.version, 11);
1813 assert_eq!(location.naming_scheme, naming_scheme);
1814 assert_eq!(location.path, naming_scheme.manifest_path(&base, 11));
1815 }
1816
1817 fn non_lexical_memory_store() -> Box<ObjectStore> {
1820 let mut object_store = ObjectStore::memory();
1821 object_store.list_is_lexically_ordered = false;
1822 Box::new(object_store)
1823 }
1824
1825 #[tokio::test]
1826 async fn test_write_version_hint() {
1827 let base = Path::from("base");
1828
1829 let lexical = ObjectStore::memory();
1831 write_version_hint(&lexical, &base, 42).await;
1832 assert_eq!(read_version_from_hint(&lexical, &base).await, None);
1833
1834 let object_store = non_lexical_memory_store();
1835 write_version_hint(&object_store, &base, 42).await;
1836 assert_eq!(read_version_from_hint(&object_store, &base).await, Some(42));
1837
1838 write_version_hint(&object_store, &base, 100).await;
1840 assert_eq!(
1841 read_version_from_hint(&object_store, &base).await,
1842 Some(100)
1843 );
1844
1845 write_version_hint(
1847 &object_store,
1848 &base,
1849 crate::format::DETACHED_VERSION_MASK | 7,
1850 )
1851 .await;
1852 assert_eq!(
1853 read_version_from_hint(&object_store, &base).await,
1854 Some(100)
1855 );
1856
1857 let hint_path = version_hint_path(&base);
1859 object_store
1860 .put(&hint_path, b"not json".as_slice())
1861 .await
1862 .unwrap();
1863 assert_eq!(read_version_from_hint(&object_store, &base).await, None);
1864 }
1865
1866 #[tokio::test]
1867 #[rstest::rstest]
1868 async fn test_read_version_hint_and_probe(
1869 #[values(ManifestNamingScheme::V1, ManifestNamingScheme::V2)]
1870 naming_scheme: ManifestNamingScheme,
1871 ) {
1872 let object_store = non_lexical_memory_store();
1873 let base = Path::from("base");
1874
1875 assert!(
1877 read_version_hint_and_probe(&object_store, &base)
1878 .await
1879 .is_none()
1880 );
1881
1882 for version in 1..=5 {
1883 object_store
1884 .put(&naming_scheme.manifest_path(&base, version), b"".as_slice())
1885 .await
1886 .unwrap();
1887 }
1888
1889 write_version_hint(&object_store, &base, 3).await;
1891 let location = read_version_hint_and_probe(&object_store, &base)
1892 .await
1893 .unwrap();
1894 assert_eq!(location.version, 5);
1895 assert_eq!(location.naming_scheme, naming_scheme);
1896
1897 write_version_hint(&object_store, &base, 5).await;
1899 let location = read_version_hint_and_probe(&object_store, &base)
1900 .await
1901 .unwrap();
1902 assert_eq!(location.version, 5);
1903
1904 write_version_hint(&object_store, &base, 10).await;
1906 assert!(
1907 read_version_hint_and_probe(&object_store, &base)
1908 .await
1909 .is_none()
1910 );
1911 }
1912
1913 #[tokio::test]
1914 async fn test_list_manifests_since_version_with_hint() {
1915 let object_store = non_lexical_memory_store();
1916 let base = Path::from("base");
1917 let scheme = ManifestNamingScheme::V2;
1918
1919 for version in 1..=10 {
1920 object_store
1921 .put(&scheme.manifest_path(&base, version), b"".as_slice())
1922 .await
1923 .unwrap();
1924 }
1925
1926 assert!(
1928 list_manifests_since_version_with_hint(&object_store, &base, 7)
1929 .await
1930 .is_none()
1931 );
1932
1933 write_version_hint(&object_store, &base, 10).await;
1935 assert!(matches!(
1936 list_manifests_since_version_with_hint(&object_store, &base, 10).await,
1937 Some(v) if v.is_empty()
1938 ));
1939
1940 let locations = list_manifests_since_version_with_hint(&object_store, &base, 7)
1943 .await
1944 .unwrap();
1945 assert_eq!(
1946 locations.iter().map(|l| l.version).collect::<Vec<_>>(),
1947 vec![10, 9, 8]
1948 );
1949
1950 write_version_hint(&object_store, &base, 8).await;
1952 let locations = list_manifests_since_version_with_hint(&object_store, &base, 7)
1953 .await
1954 .unwrap();
1955 assert_eq!(
1956 locations.iter().map(|l| l.version).collect::<Vec<_>>(),
1957 vec![10, 9, 8]
1958 );
1959
1960 write_version_hint(&object_store, &base, 20).await;
1962 assert!(
1963 list_manifests_since_version_with_hint(&object_store, &base, 7)
1964 .await
1965 .is_none()
1966 );
1967 }
1968
1969 #[tokio::test]
1970 async fn test_current_manifest_path_with_hint_non_lexical() {
1971 let object_store = non_lexical_memory_store();
1973 let base = Path::from("base");
1974 let naming_scheme = ManifestNamingScheme::V2;
1975
1976 for version in 1..=100 {
1977 object_store
1978 .put(&naming_scheme.manifest_path(&base, version), b"".as_slice())
1979 .await
1980 .unwrap();
1981 }
1982
1983 write_version_hint(&object_store, &base, 98).await;
1985 let location = current_manifest_path(&object_store, &base).await.unwrap();
1986 assert_eq!(location.version, 100);
1987 }
1988
1989 #[tokio::test]
1990 async fn test_current_manifest_path_with_stale_hint_falls_back_to_listing() {
1991 let object_store = non_lexical_memory_store();
1992 let base = Path::from("base");
1993 let naming_scheme = ManifestNamingScheme::V2;
1994
1995 object_store
1997 .put(&naming_scheme.manifest_path(&base, 5), b"".as_slice())
1998 .await
1999 .unwrap();
2000 write_version_hint(&object_store, &base, 10).await;
2001
2002 let location = current_manifest_path(&object_store, &base).await.unwrap();
2004 assert_eq!(location.version, 5);
2005 }
2006
2007 #[test]
2008 fn test_parse_detached_version() {
2009 assert_eq!(
2011 ManifestNamingScheme::parse_detached_version("d12345.manifest"),
2012 Some(12345)
2013 );
2014 assert_eq!(
2015 ManifestNamingScheme::parse_detached_version("d9223372036854775808.manifest"),
2016 Some(9223372036854775808)
2017 );
2018
2019 assert_eq!(
2021 ManifestNamingScheme::parse_detached_version("12345.manifest"),
2022 None
2023 );
2024
2025 assert_eq!(
2027 ManifestNamingScheme::parse_detached_version("18446744073709551615.manifest"),
2028 None
2029 );
2030
2031 assert_eq!(ManifestNamingScheme::parse_detached_version("d12345"), None);
2033 }
2034
2035 #[tokio::test]
2036 async fn test_list_detached_manifests() {
2037 use crate::format::DETACHED_VERSION_MASK;
2038 use futures::TryStreamExt;
2039
2040 let object_store = ObjectStore::memory();
2041 let base = Path::from("base");
2042 let versions_dir = base.clone().join(VERSIONS_DIR);
2043
2044 for version in [1, 2, 3] {
2046 let path = ManifestNamingScheme::V2.manifest_path(&base, version);
2047 object_store.put(&path, b"".as_slice()).await.unwrap();
2048 }
2049
2050 let detached_versions: Vec<u64> = vec![
2052 100 | DETACHED_VERSION_MASK,
2053 200 | DETACHED_VERSION_MASK,
2054 300 | DETACHED_VERSION_MASK,
2055 ];
2056 for version in &detached_versions {
2057 let path = versions_dir.clone().join(format!("d{}.manifest", version));
2058 object_store.put(&path, b"".as_slice()).await.unwrap();
2059 }
2060
2061 let detached_locations: Vec<ManifestLocation> =
2063 list_detached_manifests(&base, &object_store.inner)
2064 .try_collect()
2065 .await
2066 .unwrap();
2067
2068 assert_eq!(detached_locations.len(), 3);
2069 for loc in &detached_locations {
2070 assert_eq!(loc.naming_scheme, ManifestNamingScheme::V2);
2071 }
2072
2073 let mut found_versions: Vec<u64> = detached_locations.iter().map(|l| l.version).collect();
2074 found_versions.sort();
2075 let mut expected_versions = detached_versions.clone();
2076 expected_versions.sort();
2077 assert_eq!(found_versions, expected_versions);
2078 }
2079
2080 #[tokio::test]
2081 #[rstest::rstest]
2082 #[case::memory("memory://bucket-a/ds")]
2083 #[case::shared_memory("shared-memory://bucket-a/ds")]
2084 #[case::s3("s3://bucket-a/ds")]
2085 #[case::gs("gs://bucket-a/ds")]
2086 #[case::az("az://bucket-a/ds")]
2087 #[case::abfss("abfss://bucket-a/ds")]
2088 #[case::oss("oss://bucket-a/ds")]
2089 #[case::tos("tos://bucket-a/ds")]
2090 #[case::goosefs("goosefs://bucket-a/ds")]
2091 async fn test_commit_handler_from_url_conditional_put_schemes(#[case] url: &str) {
2092 let handler = commit_handler_from_url(url, &None).await.unwrap();
2097 assert_eq!(
2098 format!("{:?}", handler),
2099 "ConditionalPutCommitHandler",
2100 "{url} should route to ConditionalPutCommitHandler",
2101 );
2102 }
2103
2104 #[derive(Debug)]
2107 struct TrackingLock {
2108 released: Arc<AtomicBool>,
2109 }
2110
2111 struct TrackingLease {
2112 released: Arc<AtomicBool>,
2113 }
2114
2115 #[async_trait::async_trait]
2116 impl CommitLock for TrackingLock {
2117 type Lease = TrackingLease;
2118 async fn lock(&self, _version: u64) -> std::result::Result<Self::Lease, CommitError> {
2119 Ok(TrackingLease {
2120 released: self.released.clone(),
2121 })
2122 }
2123 }
2124
2125 #[async_trait::async_trait]
2126 impl CommitLease for TrackingLease {
2127 async fn release(&self, _success: bool) -> std::result::Result<(), CommitError> {
2128 self.released
2129 .store(true, std::sync::atomic::Ordering::SeqCst);
2130 Ok(())
2131 }
2132 }
2133
2134 #[derive(Debug)]
2138 struct HangingReleaseLock {
2139 release_calls: Arc<AtomicUsize>,
2140 released: Arc<AtomicBool>,
2141 }
2142
2143 struct HangingReleaseLease {
2144 release_calls: Arc<AtomicUsize>,
2145 released: Arc<AtomicBool>,
2146 }
2147
2148 #[async_trait::async_trait]
2149 impl CommitLock for HangingReleaseLock {
2150 type Lease = HangingReleaseLease;
2151 async fn lock(&self, _version: u64) -> std::result::Result<Self::Lease, CommitError> {
2152 Ok(HangingReleaseLease {
2153 release_calls: self.release_calls.clone(),
2154 released: self.released.clone(),
2155 })
2156 }
2157 }
2158
2159 #[async_trait::async_trait]
2160 impl CommitLease for HangingReleaseLease {
2161 async fn release(&self, _success: bool) -> std::result::Result<(), CommitError> {
2162 if self
2167 .release_calls
2168 .fetch_add(1, std::sync::atomic::Ordering::SeqCst)
2169 == 0
2170 {
2171 future::pending::<()>().await;
2172 unreachable!()
2173 }
2174 self.released
2175 .store(true, std::sync::atomic::Ordering::SeqCst);
2176 Ok(())
2177 }
2178 }
2179
2180 fn succeeding_manifest_writer<'a>(
2183 _object_store: &'a ObjectStore,
2184 _manifest: &'a mut Manifest,
2185 _indices: Option<Vec<IndexMetadata>>,
2186 _path: &'a Path,
2187 _transaction: Option<Transaction>,
2188 ) -> BoxFuture<'a, Result<WriteResult>> {
2189 Box::pin(async move { Ok(WriteResult::default()) })
2190 }
2191
2192 fn test_manifest() -> Manifest {
2193 use std::collections::HashMap;
2194
2195 use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema};
2196 use lance_core::datatypes::Schema;
2197 use lance_file::version::LanceFileVersion;
2198
2199 use crate::format::DataStorageFormat;
2200
2201 let arrow_schema = ArrowSchema::new(vec![ArrowField::new("i", DataType::Int32, false)]);
2202 Manifest::new(
2203 Schema::try_from(&arrow_schema).unwrap(),
2204 Arc::new(vec![]),
2205 DataStorageFormat::new(LanceFileVersion::Stable.resolve()),
2206 HashMap::new(),
2207 )
2208 }
2209
2210 #[tokio::test]
2211 async fn test_cos_commit_requires_custom_handler() {
2212 let handler = commit_handler_from_url("cos://bucket-a/ds", &None)
2213 .await
2214 .unwrap();
2215 assert_eq!(format!("{:?}", handler), "TencentCosCommitHandler");
2216
2217 let mut manifest = test_manifest();
2218 let error = handler
2219 .commit(
2220 &mut manifest,
2221 None,
2222 &Path::from("test"),
2223 &ObjectStore::memory(),
2224 succeeding_manifest_writer,
2225 ManifestNamingScheme::V2,
2226 None,
2227 )
2228 .await
2229 .unwrap_err();
2230 let CommitError::OtherError(error) = error else {
2231 panic!("expected a not-supported commit error");
2232 };
2233 assert!(matches!(error, Error::NotSupported { .. }));
2234 assert!(error.to_string().contains("distributed commit_lock"));
2235 }
2236
2237 fn hanging_manifest_writer<'a>(
2239 _object_store: &'a ObjectStore,
2240 _manifest: &'a mut Manifest,
2241 _indices: Option<Vec<IndexMetadata>>,
2242 _path: &'a Path,
2243 _transaction: Option<Transaction>,
2244 ) -> BoxFuture<'a, Result<WriteResult>> {
2245 Box::pin(async move {
2246 future::pending::<()>().await;
2247 unreachable!()
2248 })
2249 }
2250
2251 #[tokio::test]
2254 async fn test_commit_lock_released_on_cancellation() {
2255 use std::sync::atomic::Ordering;
2256 use std::time::Duration;
2257
2258 let released = Arc::new(AtomicBool::new(false));
2259 let lock = TrackingLock {
2260 released: released.clone(),
2261 };
2262
2263 let object_store = ObjectStore::memory();
2264 let base_path = Path::from("test");
2265 let mut manifest = test_manifest();
2266
2267 let commit_fut = lock.commit(
2270 &mut manifest,
2271 None,
2272 &base_path,
2273 &object_store,
2274 hanging_manifest_writer,
2275 ManifestNamingScheme::V2,
2276 None,
2277 );
2278 let timed_out = tokio::time::timeout(Duration::from_millis(50), commit_fut).await;
2279 assert!(timed_out.is_err(), "commit should not have completed");
2280
2281 for _ in 0..100 {
2283 if released.load(Ordering::SeqCst) {
2284 break;
2285 }
2286 tokio::time::sleep(Duration::from_millis(10)).await;
2287 }
2288 assert!(
2289 released.load(Ordering::SeqCst),
2290 "lock must be released after the commit future is cancelled"
2291 );
2292 }
2293
2294 #[tokio::test]
2298 async fn test_commit_lock_released_on_cancellation_during_release() {
2299 use std::sync::atomic::Ordering;
2300 use std::time::Duration;
2301
2302 let release_calls = Arc::new(AtomicUsize::new(0));
2303 let released = Arc::new(AtomicBool::new(false));
2304 let lock = HangingReleaseLock {
2305 release_calls: release_calls.clone(),
2306 released: released.clone(),
2307 };
2308
2309 let object_store = ObjectStore::memory();
2310 let base_path = Path::from("test");
2311 let mut manifest = test_manifest();
2312
2313 let commit_fut = lock.commit(
2316 &mut manifest,
2317 None,
2318 &base_path,
2319 &object_store,
2320 succeeding_manifest_writer,
2321 ManifestNamingScheme::V2,
2322 None,
2323 );
2324 let timed_out = tokio::time::timeout(Duration::from_millis(50), commit_fut).await;
2325 assert!(timed_out.is_err(), "commit should not have completed");
2326
2327 for _ in 0..100 {
2330 if released.load(Ordering::SeqCst) {
2331 break;
2332 }
2333 tokio::time::sleep(Duration::from_millis(10)).await;
2334 }
2335 assert!(
2336 released.load(Ordering::SeqCst),
2337 "lock must be released even when cancelled during the explicit release"
2338 );
2339 assert_eq!(
2340 release_calls.load(Ordering::SeqCst),
2341 2,
2342 "expected the hung explicit release plus one best-effort drop release"
2343 );
2344 }
2345}