Skip to main content

lance_table/io/
commit.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Trait for commit implementations.
5//!
6//! In Lance, a transaction is committed by writing the next manifest file.
7//! However, care should be taken to ensure that the manifest file is written
8//! only once, even if there are concurrent writers. Different stores have
9//! different abilities to handle concurrent writes, so a trait is provided
10//! to allow for different implementations.
11//!
12//! The trait [CommitHandler] can be implemented to provide different commit
13//! strategies. The default implementation for most object stores is
14//! [RenameCommitHandler], which writes the manifest to a temporary path, then
15//! renames the temporary path to the final path if no object already exists
16//! at the final path. This is an atomic operation in most object stores, but
17//! not in AWS S3. So for AWS S3, the default commit handler is
18//! [UnsafeCommitHandler], which writes the manifest to the final path without
19//! any checks.
20//!
21//! When providing your own commit handler, most often you are implementing in
22//! terms of a lock. The trait [CommitLock] can be implemented as a simpler
23//! alternative to [CommitHandler].
24
25use 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";
73/// File name for the JSON version hint file, stored under `_versions/`.
74///
75/// The file contains `{"version":N}` where `N` is the latest committed version
76/// at the time of writing. It enables O(1)/O(k) latest-version lookup via HEAD
77/// requests on object stores where listing is not lexicographically ordered
78/// (e.g. S3 Express, local filesystem) instead of an O(n) listing.
79const VERSION_HINT_FILE: &str = "latest_version_hint.json";
80
81/// How manifest files should be named.
82#[derive(Clone, Copy, Debug, PartialEq, Eq)]
83pub enum ManifestNamingScheme {
84    /// `_versions/{version}.manifest`
85    V1,
86    /// `_manifests/{u64::MAX - version}.manifest`
87    ///
88    /// Zero-padded and reversed for O(1) lookup of latest version on object stores.
89    V2,
90}
91
92impl ManifestNamingScheme {
93    pub fn manifest_path(&self, base: &Path, version: u64) -> Path {
94        if is_detached_version(version) {
95            // Detached versions should never show up first in a list operation which
96            // means it needs to come lexicographically after all attached manifest
97            // files and so we add the prefix `d`.  There is no need to invert the
98            // version number since detached versions are not part of the version
99            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            // Detached versions will fail the `parse` step, which is ok.
118            .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    /// Parse a detached version from a filename like `d123456.manifest`.
126    ///
127    /// Returns the full version number with the detached mask bit set.
128    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            // Currently, detached versions must imply V2
141            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        // We shouldn't have to worry about detached versions here since there is no
157        // such thing as "detached" and "staged" at the same time.
158        if filename.chars().nth(20) == Some('.') {
159            Self::V2
160        } else {
161            Self::V1
162        }
163    }
164}
165
166/// Migrate all V1 manifests to V2 naming scheme.
167///
168/// This function will rename all V1 manifests to V2 naming scheme.
169///
170/// This function is idempotent, and can be run multiple times without
171/// changing the state of the object store.
172///
173/// However, it should not be run while other concurrent operations are happening.
174/// And it should also run until completion before resuming other operations.
175pub 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
199/// Function that writes the manifest to the object store.
200///
201/// Returns the size of the written manifest.
202pub 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
210/// Canonical manifest writer; its function item type exactly matches `ManifestWriter`.
211/// Rationale: keep a crate-local writer implementation so call sites can pass this function
212/// directly without non-primitive casts or lifetime coercions.
213pub 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    /// The version the manifest corresponds to.
235    pub version: u64,
236    /// Path of the manifest file, relative to the table root.
237    pub path: Path,
238    /// Size, in bytes, of the manifest file. If it is not known, this field should be `None`.
239    pub size: Option<u64>,
240    /// Naming scheme of the manifest file.
241    pub naming_scheme: ManifestNamingScheme,
242    /// Optional opaque object generation token observed at `path`.
243    ///
244    /// An ETag is not necessarily a content checksum and may change when an
245    /// object is rewritten with identical bytes. In particular, S3 Express
246    /// returns an object-specific opaque value. Callers must not treat it as a
247    /// content checksum, logical manifest identity, or dataset-incarnation
248    /// identity. The generic
249    /// [`ExternalManifestStore`](crate::io::commit::external_manifest::ExternalManifestStore)
250    /// workflow therefore neither persists nor validates it: COPY and external
251    /// index publication are not atomic, so an otherwise correct equivalent
252    /// materialization can make a stored token stale before it is published.
253    ///
254    /// When present, the token still distinguishes the physical object
255    /// generation observed by this caller and can prevent reuse of an older
256    /// cached Dataset at the same URI and version. Conversely, `None` must not
257    /// be interpreted as proof that two observations belong to the same dataset
258    /// incarnation.
259    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
284/// Get the latest manifest path.
285///
286/// - Local filesystem: a single directory read.
287/// - Stores where listing is not lexicographically ordered (e.g. S3 Express):
288///   the version hint (read the hint file, then probe higher versions with
289///   HEADs), falling back to a listing if the hint is missing or stale. A full
290///   listing on these stores is O(n) in the number of versions.
291/// - Lexicographically ordered stores (e.g. S3 Standard, GCS): the listing
292///   already resolves the latest version in roughly one request.
293async 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/// JSON body of the version hint file: `{"version":N}`.
311#[derive(serde::Serialize, serde::Deserialize)]
312struct VersionHint {
313    version: u64,
314}
315
316/// Set `LANCE_USE_VERSION_HINT=0` (or `false`) to globally disable the version
317/// hint — writers stop emitting the hint file and readers stop consulting it,
318/// falling back to plain listing. Intended as a benchmark/escape-hatch knob;
319/// the hint is on by default.
320const 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
333/// Whether this object store benefits from a version hint.
334///
335/// On stores where listing is lexicographically ordered (S3 Standard, GCS,
336/// Azure, ...) the latest version is already resolved in roughly one request,
337/// so the hint would only add a write per commit for nothing. We write (and
338/// read) it only on stores where listing is not lexicographically ordered —
339/// S3 Express and the local filesystem. Can be force-disabled with the
340/// `LANCE_USE_VERSION_HINT=0` environment variable.
341pub fn uses_version_hint(object_store: &ObjectStore) -> bool {
342    version_hint_globally_enabled() && !object_store.list_is_lexically_ordered
343}
344
345/// Path to the JSON version hint file for a dataset.
346fn version_hint_path(base: &Path) -> Path {
347    base.clone().join(VERSIONS_DIR).join(VERSION_HINT_FILE)
348}
349
350/// Write the version hint file after a successful commit.
351///
352/// The hint is stored as JSON: `{"version":N}`. This write is best-effort —
353/// failures are logged and ignored, since the hint only accelerates reads and
354/// never affects correctness (readers verify the hinted version and probe
355/// upward from there). It is a no-op for detached versions and for stores that
356/// do not benefit from a hint (see [`uses_version_hint`]).
357pub 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
368/// Read the latest version from the hint file, or `None` if it does not exist
369/// or cannot be parsed.
370async 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
382/// Read the version hint and probe upward to find the true latest manifest.
383///
384/// Returns `None` if the hint file is missing, the hinted version no longer
385/// exists, or any error occurred — callers should fall back to listing.
386async 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    // `probed` is non-empty and its last entry is the highest version found.
396    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
406/// Maximum version gap between the hint and the read version for which we use
407/// the hint-based parallel-HEAD path; beyond this a single (paginated) listing
408/// is cheaper, so callers fall back to it.
409const MAX_HINT_PROBE_GAP: u64 = 1000;
410
411/// Probe `from_version`, then `from_version + 1`, `+ 2`, ... with HEAD requests
412/// until one is not found.
413///
414/// Assumes attached versions are contiguous above `from_version` (true in
415/// practice: every commit increments by one, and cleanup only removes *old*
416/// versions, never ones newer than the latest). A `NotFound` therefore marks
417/// the end of the history.
418///
419/// - `Ok(Some((true_latest_version, naming_scheme, [(version, meta), ...])))`:
420///   the vec covers every version from `from_version` through the true latest
421///   in ascending order.
422/// - `Ok(None)`: `from_version` itself does not exist (a `NotFound` for both
423///   naming schemes) — i.e. the hint pointed past the end.
424/// - `Err(_)`: a transient object-store error was hit, so the probed range may
425///   be incomplete; callers should fall back to a full listing rather than
426///   trust a possibly-stale result.
427async 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    // Newer datasets use V2; fall back to V1 if the V2 path is not found.
439    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            // NotFound means we found the latest version.
475            Err(ObjectStoreError::NotFound { .. }) => break,
476            // A transient error means a newer version might exist that we
477            // failed to observe — surface it so callers fall back to listing.
478            Err(e) => return Err(e.into()),
479        }
480    }
481    Ok(Some((version, scheme, probed)))
482}
483
484/// List manifest locations with version `> since_version` using the version
485/// hint, in descending order of version.
486///
487/// Returns `None` if the hint is missing or stale enough that this is not
488/// usable — callers should fall back to a full listing. `Some(vec![])` is the
489/// fast path where the hint confirms there are no new versions.
490async 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    // A reader that is very far behind is cheaper to serve with one paginated
498    // listing than with thousands of HEADs.
499    if hint_version.saturating_sub(since_version) > MAX_HINT_PROBE_GAP {
500        return None;
501    }
502
503    // If the hint is not newer than the read version, the only versions that
504    // could exist are right above it; otherwise start at the hint.
505    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        // Nothing at `probe_from`. If we were probing from the hint, the hint
514        // is stale — bail to a full listing. If we were probing from
515        // `since_version + 1`, there are simply no new versions.
516        Ok(None) if hint_version > since_version => return None,
517        Ok(None) => return Some(Vec::new()),
518        // Transient error: don't trust the hint path, fall back to listing.
519        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    // Fill the gap between `since_version` and the hint with HEADs (the probe
535    // above already covered `hint_version` and up). The range is contiguous, so
536    // any error here (including a `NotFound`) means we can't trust the hint path
537    // — fall back to a full listing.
538    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
565/// Resolve the latest manifest by listing the versions directory.
566async 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            // Only include if we can parse a version (skip detached versions)
576            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        // If the first valid manifest we see is V2, we can assume that we are using
589        // V2 naming scheme for all manifests.
590        (Some((scheme @ ManifestNamingScheme::V2, meta)), true) => {
591            let version = scheme
592                .parse_version(meta.location.filename().unwrap())
593                .unwrap();
594
595            // Sanity check: verify at least for the first 1k files that they are all V2
596            // and that the version numbers are decreasing. We use the first 1k because
597            // this is the typical size of an object store list endpoint response page.
598            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        // If the list is not lexically ordered, we need to iterate all manifests
628        // to find the latest version. This works for both V1 and V2 schemes.
629        (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
666// This is an optimized function that searches for the latest manifest. In
667// object_store, list operations lookup metadata for each file listed. This
668// method only gets the metadata for the found latest manifest.
669fn 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            // Need to ignore temporary files, such as
684            // .tmp_7.manifest_9c100374-3298-4537-afc6-f5ee7913666d
685            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
745/// Convert object metadata to ManifestLocation for detached manifests.
746fn 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
760/// List all detached manifest files in the versions directory.
761pub 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/// Handle commits that prevent conflicting writes.
786///
787/// Commit implementations ensure that if there are multiple concurrent writers
788/// attempting to write the next version of a table, only one will win. In order
789/// to work, all writers must use the same commit handler type.
790/// This trait is also responsible for resolving where the manifests live.
791///
792// TODO: pub(crate)
793#[async_trait::async_trait]
794#[allow(clippy::too_many_arguments)]
795pub trait CommitHandler: Debug + Send + Sync {
796    /// Whether a not-found result from [`Self::resolve_version_location`] is
797    /// definitive immediately after a commit attempt.
798    ///
799    /// Handlers backed by an eventually consistent or external source of
800    /// truth should keep the conservative default. This prevents callers from
801    /// deleting files that a newly committed manifest may reference while the
802    /// manifest is not yet visible through the resolver.
803    fn is_version_not_found_definitive(&self) -> bool {
804        false
805    }
806
807    /// Whether an error should still be returned after readback proves that
808    /// the manifest from the current commit attempt landed.
809    ///
810    /// The conservative default preserves errors from custom handlers. Built-in
811    /// object-store handlers override this because their commit errors may be
812    /// ambiguous transport failures whose successful outcome is authoritative.
813    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    /// Check whether an attached manifest version exists without loading it.
835    ///
836    /// The default implementation probes the deterministic manifest path for
837    /// the given naming scheme. Commit handlers with an external source of
838    /// truth should override this method.
839    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    /// List detached manifest locations.
855    ///
856    /// Returns a stream of detached manifest locations in arbitrary order.
857    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    /// If `sorted_descending` is `true`, the stream will yield manifests in descending
866    /// order of version. When the object store has a lexicographically
867    /// ordered list and the naming scheme is V2, this will use an optimized
868    /// list operation. Otherwise, it will list all manifests and sort them
869    /// in memory. When `sorted_descending` is `false`, the stream will yield manifests
870    /// in arbitrary order.
871    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 the object store supports lexicographically ordered lists and
892        // the naming scheme is V2, we can use an optimized list operation.
893        if object_store.list_is_lexically_ordered {
894            // We don't know the naming scheme until we see the first manifest.
895            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                    // If we get an error or no manifests are found, we default
901                    // to V2 naming scheme, since it doesn't matter.
902                    Some(Err(_)) => ManifestNamingScheme::V2,
903                    None => ManifestNamingScheme::V2,
904                };
905
906                if naming_scheme == ManifestNamingScheme::V2 {
907                    // If the first manifest is V2, we can use the optimized list operation.
908                    Ok(Either::Left(peekable))
909                } else {
910                    sort_stream(peekable).await.map(Either::Right)
911                }
912            })
913            .try_flatten()
914            .boxed()
915        } else {
916            // If the object store does not support lexicographically ordered lists,
917            // we need to sort the manifests in memory. Systems where this isn't
918            // supported (local fs, S3 express) are typically fast enough
919            // that this is not a problem.
920            futures::stream::once(sort_stream(underlying_stream))
921                .try_flatten()
922                .boxed()
923        }
924    }
925
926    /// List manifest locations with version `> since_version`, in descending
927    /// order of version.
928    ///
929    /// On lexically-ordered stores this is the standard listing with early
930    /// termination. On non-lexically-ordered stores (e.g. S3 Express) it uses
931    /// the version hint to avoid an O(n) listing, falling back to a full
932    /// listing if the hint is missing or stale.
933    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    /// Commit a manifest.
972    ///
973    /// This function should return an [CommitError::CommitConflict] if another
974    /// transaction has already been committed to the path.
975    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    /// Delete the recorded manifest information for a dataset at the base_path
987    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            // Detached versions are not supported with V1 naming scheme.  If we need
1001            // to support in the future we could use a different prefix (e.g. 'x' or something)
1002            naming_scheme: ManifestNamingScheme::V2,
1003            // Both V1 and V2 should give the same path for detached versions
1004            path: ManifestNamingScheme::V2.manifest_path(base_path, version),
1005            size: None,
1006            e_tag: None,
1007        });
1008    }
1009
1010    // try V2, fallback to V1.
1011    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            // fallback to V1
1023            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/// Adapt an object_store credentials into AWS SDK creds
1036#[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, //  10 min
1062                        ))
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        // caching should be handled by passed AwsCredentialProvider
1090        .identity_cache(IdentityCache::no_cache())
1091        // Be more resilient to transient network issues.
1092        // 5 attempts = 1 initial + 4 retries with exponential backoff.
1093        .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    // This looks unused if dynamodb feature disabled
1106    #[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            // On Windows, the drive is parsed as a scheme
1117            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            // Get accessor from the options
1168            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                    &region,
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/// Errors that can occur when committing a manifest.
1207#[derive(Debug)]
1208pub enum CommitError {
1209    /// Another transaction has already been written to the path
1210    CommitConflict,
1211    /// Something else went wrong
1212    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
1230/// Whether we have issued a warning about using the unsafe commit handler.
1231static WARNED_ON_UNSAFE_COMMIT: AtomicBool = AtomicBool::new(false);
1232
1233/// A naive commit implementation that does not prevent conflicting writes.
1234///
1235/// This will log a warning the first time it is used.
1236pub 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        // Log a one-time warning
1260        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/// A commit implementation that uses a lock to prevent conflicting writes.
1291#[async_trait::async_trait]
1292pub trait CommitLock: Debug {
1293    type Lease: CommitLease;
1294
1295    /// Attempt to lock the table for the given version.
1296    ///
1297    /// If it is already locked by another transaction, wait until it is unlocked.
1298    /// Once it is unlocked, return [CommitError::CommitConflict] if the version
1299    /// has already been committed. Otherwise, return the lock.
1300    ///
1301    /// To prevent poisoned locks, it's recommended to set a timeout on the lock
1302    /// of at least 30 seconds.
1303    ///
1304    /// It is not required that the lock tracks the version. It is provided in
1305    /// case the locking is handled by a catalog service that needs to know the
1306    /// current version of the table.
1307    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    /// Return the lease, indicating whether the commit was successful.
1313    ///
1314    /// Implementations should tolerate being called more than once: if a commit
1315    /// is cancelled (e.g. by a timeout) while `release` is in flight, a
1316    /// best-effort `release(false)` may be issued afterwards from the drop path.
1317    async fn release(&self, success: bool) -> std::result::Result<(), CommitError>;
1318}
1319
1320/// Guards a [CommitLease] so the lock is released even if the commit future is
1321/// dropped (e.g. cancelled by a commit timeout) before reaching an explicit
1322/// release.
1323///
1324/// [CommitLease::release] is async and cannot be awaited from `Drop`, so on the
1325/// drop path we spawn a best-effort background task that releases the lock with
1326/// `success = false`. Without this, a cancelled commit would leak the lock until
1327/// the lease's own TTL expired, blocking other writers in the meantime.
1328struct 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    /// Explicitly release the lease, consuming the guard so `Drop` is a no-op.
1338    async fn release(mut self, success: bool) -> std::result::Result<(), CommitError> {
1339        // Keep the lease inside the guard across the await so that, if this
1340        // future is cancelled mid-release (e.g. the release call itself hangs
1341        // and the commit timeout fires), `Drop` still issues a best-effort
1342        // release. Only clear it once the release has fully completed.
1343        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            // The guard was dropped without an explicit release, meaning the
1359            // commit future was cancelled while holding the lock. We can't await
1360            // in `Drop`, so spawn a best-effort release. If there is no runtime,
1361            // leave the lease for its TTL to reclaim.
1362            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        // Hold the lease in a guard so the lock is released even if this future
1392        // is cancelled before we reach an explicit release below. The explicit
1393        // releases are still preferred since they report the correct success
1394        // flag and surface release errors; the guard only covers cancellation.
1395        let lease = LeaseGuard::new(self.lock(manifest.version).await?);
1396
1397        // Head the location and make sure it's not already committed
1398        match object_store.inner.head(&path).await {
1399            Ok(_) => {
1400                // The path already exists, so it's already committed
1401                // Release the lock
1402                lease.release(false).await?;
1403
1404                return Err(CommitError::CommitConflict);
1405            }
1406            Err(ObjectStoreError::NotFound { .. }) => {}
1407            Err(e) => {
1408                // Something else went wrong
1409                // Release the lock
1410                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        // Release the lock
1418        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
1471/// A commit implementation that uses a temporary path and renames the object.
1472///
1473/// This only works for object stores that support atomic rename if not exist.
1474pub 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        // Create a temporary object, then use `rename_if_not_exists` to commit.
1497        // If failed, clean up the temporary object.
1498
1499        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                // Successfully committed
1511                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, // Re-name can change e-tag.
1518                })
1519            }
1520            Err(ObjectStoreError::AlreadyExists { .. }) => {
1521                // Another transaction has already been committed
1522                // Attempt to clean up temporary object, but ignore errors if we can't
1523                let _ = object_store.delete(&tmp_path).await;
1524
1525                return Err(CommitError::CommitConflict);
1526            }
1527            Err(e) => {
1528                // Something else went wrong
1529                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
1613/// A read-capable handler that prevents unsafe default commits to Tencent COS.
1614///
1615/// COS silently ignores its put-if-not-exists header on buckets that have ever
1616/// had versioning enabled. Since that bucket history cannot be inferred from
1617/// the URI or storage options, using [`ConditionalPutCommitHandler`] here can
1618/// let concurrent writers overwrite the same manifest without reporting a
1619/// conflict.
1620struct 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    // TODO: add isolation_level
1657}
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        // Write two v1 files and one v1
1728        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        // Write 12 manifest files, latest first
1773        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        // Use memory store for both cases to avoid local FS special codepath.
1798        // Modify list_is_lexically_ordered to simulate different object stores.
1799        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        // Write 12 manifest files in non-sequential order
1805        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    /// A memory store that reports `list_is_lexically_ordered == false`, like
1818    /// S3 Express, so the version-hint paths are exercised.
1819    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        // No hint is written on lexically-ordered stores (it would not be read).
1830        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        // A later commit overwrites the hint.
1839        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        // Detached versions are never written to the hint.
1846        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        // A corrupt / non-JSON hint file is treated as missing.
1858        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        // No hint file yet.
1876        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        // Stale hint: should probe forward and find version 5.
1890        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        // Up-to-date hint: returns version 5 directly.
1898        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        // Hint points past the latest version: not usable.
1905        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        // No hint yet -> not usable, caller must fall back.
1927        assert!(
1928            list_manifests_since_version_with_hint(&object_store, &base, 7)
1929                .await
1930                .is_none()
1931        );
1932
1933        // Hint exactly at the read version -> fast path, nothing new.
1934        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        // Hint ahead of the read version, with a gap to fill (8, 9) plus probing
1941        // from the hint (10). Results are descending by version.
1942        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        // Slightly stale hint (points at 8) still probes up to the true latest.
1951        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        // Hint points past the latest -> not usable, caller falls back.
1961        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        // Simulate S3 Express (non-lexically ordered list) with many versions.
1972        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        // Slightly stale hint: probing from 98 still resolves the true latest.
1984        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        // Only version 5 exists, but the hint claims version 10.
1996        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        // The stale hint is ignored; listing finds version 5.
2003        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        // Valid detached version filenames
2010        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        // Invalid: not starting with 'd' prefix
2020        assert_eq!(
2021            ManifestNamingScheme::parse_detached_version("12345.manifest"),
2022            None
2023        );
2024
2025        // Invalid: regular V2 manifest
2026        assert_eq!(
2027            ManifestNamingScheme::parse_detached_version("18446744073709551615.manifest"),
2028            None
2029        );
2030
2031        // Invalid: no extension
2032        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        // Create some regular manifests
2045        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        // Create some detached manifests
2051        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        // List detached manifests
2062        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        // Every scheme whose store supports atomic put-if-not-exists must
2093        // route to ConditionalPutCommitHandler — otherwise concurrent writers
2094        // fall through to UnsafeCommitHandler and silently clobber each
2095        // other's manifests.
2096        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    /// A [CommitLock] whose lease records whether it was released, so we can
2105    /// assert the lock does not leak when the commit future is cancelled.
2106    #[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    /// A [CommitLock] whose lease hangs on its first `release` call but completes
2135    /// on subsequent ones, so we can assert the drop-path best-effort release
2136    /// fires when a commit is cancelled *during* the explicit release.
2137    #[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            // The first release (the explicit one) hangs, simulating a release
2163            // call that stalls long enough for the commit timeout to fire. The
2164            // best-effort release issued from `Drop` is the second call and
2165            // succeeds.
2166            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    /// A manifest writer that succeeds immediately, so the commit reaches the
2181    /// explicit lease release.
2182    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    /// A manifest writer that never completes, simulating a hung object store.
2238    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    /// Cancelling a commit (as a commit timeout does) while the lock is held must
2252    /// still release the lock; otherwise it leaks until the lease's TTL expires.
2253    #[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        // The commit will hang on the manifest writer while holding the lock.
2268        // Cancel it the same way a commit timeout would: drop the future.
2269        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        // The drop guard releases the lock on a background task; wait for it.
2282        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    /// Cancelling a commit *during* the explicit lease release (e.g. the release
2295    /// call itself hangs and the commit timeout fires) must still release the
2296    /// lock via the drop-path best-effort release.
2297    #[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        // The manifest writer succeeds, so the commit reaches the explicit
2314        // release, which hangs. Cancel it the same way a commit timeout would.
2315        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        // The drop guard issues a best-effort release on a background task; wait
2328        // for it. This is the second release call (the first one hung).
2329        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}