moraine 0.6.4

use std::sync::Arc;

use object_store::memory::InMemory;

use crate::{
    catalog::{Catalog, CatalogOptions, TableId, projection},
    store::{
        key::{InlineKey, InlineOperation, Key},
        proto, value,
    },
    transaction::staged::{Cell, RowOperation, StagedTransaction, TableKind},
};

async fn open() -> Catalog {
    Catalog::open(Arc::new(InMemory::new()), CatalogOptions::default())
        .await
        .unwrap()
}

fn snapshot_row(id: u64) -> Vec<Cell> {
    vec![
        Cell::U64(id),
        Cell::I64(1),
        Cell::U64(0),
        Cell::U64(1),
        Cell::U64(0),
    ]
}

fn snapshot_changes_row(id: u64, changes_made: &str) -> Vec<Cell> {
    vec![
        Cell::U64(id),
        Cell::Str(changes_made.to_string()),
        Cell::Null,
        Cell::Null,
        Cell::Null,
    ]
}

/// Commits one staged batch minting `snapshot`, carrying `ops`.
async fn commit_staged(catalog: &Catalog, snapshot: u64, changes: &str, ops: Vec<RowOperation>) {
    let db_tx = catalog.begin_write_tx().await.unwrap();
    let mut tx = StagedTransaction::begin_detached_on(catalog, db_tx);
    for op in ops {
        tx.stage(op);
    }
    tx.stage(RowOperation::Insert {
        table: TableKind::Snapshot,
        cells: snapshot_row(snapshot),
    });
    tx.stage(RowOperation::Insert {
        table: TableKind::SnapshotChanges,
        cells: snapshot_changes_row(snapshot, changes),
    });
    tx.commit().await.unwrap();
}

/// The first `recent_rows` walk verifies the chunk directory and remembers
/// it; the next serves from it. The proof is divergence: a chunk planted
/// without a locator is visible only to a body scan, so a read that misses
/// it read the directory.
#[tokio::test]
async fn recent_rows_serves_from_a_directory_verified_complete() {
    let catalog = open().await;

    commit_staged(
        &catalog,
        1,
        "inlined_insert:1",
        vec![
            RowOperation::InlineSchema {
                table_id: 1,
                schema_version: 0,
                arrow_schema: b"schema".to_vec(),
            },
            RowOperation::InlineInsert {
                table_id: 1,
                schema_version: 0,
                begin_snapshot: 1,
                row_id_start: 0,
                row_count: 2,
                arrow_body: b"chunk".to_vec(),
            },
        ],
    )
    .await;

    let rows = catalog.recent_rows(TableId::new(1)).await.unwrap();
    assert_eq!(rows.len(), 2);
    assert!(
        projection::inline_directory_complete(catalog.projections(), 1),
        "the first walk must verify and remember a complete directory"
    );

    // The divergence: a chunk with no locator, behind the memo's back.
    let tx = catalog.begin_write_tx().await.unwrap();
    tx.put(
        Key::Inline(InlineKey::Live(InlineOperation::Insert {
            table_id: 1,
            schema_version: 0,
            begin_snapshot: 1,
            chunk_seq: 5,
        }))
        .encode(),
        value::encode_value(&proto::InlineChunkValue {
            body: b"rogue".to_vec().into(),
            row_id_start: 10,
            row_count: 2,
            data_file_id: None,
        }),
    )
    .unwrap();
    tx.commit().await.unwrap();

    let rows = catalog.recent_rows(TableId::new(1)).await.unwrap();
    let row_ids: Vec<u64> = rows.iter().map(|row| row.row_id).collect();
    assert_eq!(
        row_ids,
        [0, 1],
        "a directory-served read must not see a chunk only a body scan finds"
    );
    assert_eq!(rows[0].chunk_body.as_slice(), b"chunk");
    assert_eq!(rows[0].arrow_schema.as_slice(), b"schema");

    catalog.close().await.unwrap();
}

/// Two live chunks of one table can end at the same row id: an UPDATE
/// re-inlines a row that another chunk's range still ends on. A locator is
/// keyed by its range end alone, so the second chunk overwrites the first's
/// locator and takes the first's other rows out of every directory-served
/// read — and both completeness checks compare *sets* of ends, so the
/// truncated directory is declared complete and never healed. Regression.
#[tokio::test]
async fn two_chunks_ending_at_one_row_id_keep_both_locators() {
    let catalog = open().await;

    commit_staged(
        &catalog,
        1,
        "inlined_insert:1",
        vec![
            RowOperation::InlineSchema {
                table_id: 1,
                schema_version: 0,
                arrow_schema: b"schema".to_vec(),
            },
            RowOperation::InlineInsert {
                table_id: 1,
                schema_version: 0,
                begin_snapshot: 1,
                row_id_start: 5,
                row_count: 3,
                arrow_body: b"chunk-a".to_vec(),
            },
        ],
    )
    .await;

    // The re-inlined row 7: its own chunk, ending where chunk A ends.
    commit_staged(
        &catalog,
        2,
        "inlined_insert:1",
        vec![RowOperation::InlineInsert {
            table_id: 1,
            schema_version: 0,
            begin_snapshot: 2,
            row_id_start: 7,
            row_count: 1,
            arrow_body: b"chunk-b".to_vec(),
        }],
    )
    .await;

    // The first walk verifies the directory and remembers it; the next
    // serves from it, so the two must agree.
    let walked: Vec<u64> = catalog
        .recent_rows(TableId::new(1))
        .await
        .unwrap()
        .iter()
        .map(|row| row.row_id)
        .collect();
    let served: Vec<u64> = catalog
        .recent_rows(TableId::new(1))
        .await
        .unwrap()
        .iter()
        .map(|row| row.row_id)
        .collect();

    assert_eq!(
        walked, served,
        "a directory-served read lost rows the body scan found"
    );
    assert!(
        served.contains(&5) && served.contains(&6),
        "a colliding locator took chunk A's other rows out: {served:?}"
    );
    // The point of keying a locator by chunk identity: both chunks are
    // named, so the directory is complete rather than merely declining to
    // claim completeness.
    assert!(
        projection::inline_directory_complete(catalog.projections(), 1),
        "both chunks must keep their own locator"
    );

    catalog.close().await.unwrap();
}

/// A directory-served `recent_rows` point-reads only the chunks its live
/// rows reference: a chunk whose rows are all tombstoned is never fetched.
/// The proof deletes that chunk's body out from under its locator — a
/// fetch would report corruption, so a clean read never issued one.
#[tokio::test]
async fn recent_rows_reads_only_the_chunks_live_rows_reference() {
    let catalog = open().await;

    commit_staged(
        &catalog,
        1,
        "inlined_insert:1",
        vec![
            RowOperation::InlineSchema {
                table_id: 1,
                schema_version: 0,
                arrow_schema: b"schema".to_vec(),
            },
            RowOperation::InlineInsert {
                table_id: 1,
                schema_version: 0,
                begin_snapshot: 1,
                row_id_start: 0,
                row_count: 2,
                arrow_body: b"chunk-a".to_vec(),
            },
        ],
    )
    .await;
    commit_staged(
        &catalog,
        2,
        "inlined_insert:1",
        vec![RowOperation::InlineInsert {
            table_id: 1,
            schema_version: 0,
            begin_snapshot: 2,
            row_id_start: 2,
            row_count: 2,
            arrow_body: b"chunk-b".to_vec(),
        }],
    )
    .await;
    commit_staged(
        &catalog,
        3,
        "inlined_delete:1",
        vec![
            RowOperation::InlineInlineDelete {
                table_id: 1,
                row_id: 0,
                end_snapshot: 3,
            },
            RowOperation::InlineInlineDelete {
                table_id: 1,
                row_id: 1,
                end_snapshot: 3,
            },
        ],
    )
    .await;

    let rows = catalog.recent_rows(TableId::new(1)).await.unwrap();
    assert_eq!(rows.len(), 2, "chunk-a's rows are tombstoned");
    assert!(
        projection::inline_directory_complete(catalog.projections(), 1),
        "the first walk must verify and remember a complete directory"
    );

    // Remove the dead chunk's body, leaving its locator: only a read that
    // skips unreferenced chunks survives this.
    let tx = catalog.begin_write_tx().await.unwrap();
    tx.delete(
        Key::Inline(InlineKey::Live(InlineOperation::Insert {
            table_id: 1,
            schema_version: 0,
            begin_snapshot: 1,
            chunk_seq: 0,
        }))
        .encode(),
    )
    .unwrap();
    tx.commit().await.unwrap();

    let rows = catalog.recent_rows(TableId::new(1)).await.unwrap();
    let row_ids: Vec<u64> = rows.iter().map(|row| row.row_id).collect();
    assert_eq!(row_ids, [2, 3]);
    assert_eq!(rows[0].chunk_body.as_slice(), b"chunk-b");

    catalog.close().await.unwrap();
}

#[tokio::test]
async fn index_queries_reuse_views_and_observe_later_commits() {
    use std::{cell::Cell, ops::Bound};

    use crate::{ColumnDef, IndexDef};

    let catalog = open().await;
    let ids = Cell::new(None);
    catalog
        .commit(|tx| {
            let schema = tx.create_schema("s")?;
            let table = tx.create_table(
                schema,
                "t",
                &[ColumnDef {
                    name: "value".into(),
                    column_type: "BIGINT".into(),
                    nulls_allowed: true,
                    default_value: None,
                    children: Vec::new(),
                }],
            )?;
            let column = tx.columns_of(table)[0].id;
            let index = tx.create_index(
                table,
                &IndexDef {
                    name: "by_value".into(),
                    columns: vec![column],
                    unique: true,
                },
                &[],
            )?;
            ids.set(Some((table, index)));
            Ok(())
        })
        .await
        .unwrap();
    let (table, index) = ids.get().unwrap();

    for operation in 0..4 {
        projection::invalidate_current_state(catalog.projections());
        assert!(catalog.writer_head_view().is_none());
        for repeated in [false, true] {
            let before = catalog.writer_head_view();
            let rows = match operation {
                0 => catalog
                    .index_lookup(
                        table,
                        index,
                        &[crate::IndexKeyValue::Int {
                            value: 1,
                            width: crate::IntWidth::I64,
                        }],
                    )
                    .await
                    .unwrap(),
                1 => catalog.index_lookup_many(table, index, &[]).await.unwrap(),
                2 => catalog
                    .index_range(table, index, Bound::Unbounded, Bound::Unbounded, false)
                    .await
                    .unwrap(),
                _ => catalog
                    .index_nulls(table, index, vec![None], false)
                    .await
                    .unwrap(),
            };
            assert!(rows.is_empty());
            let cached = catalog
                .writer_head_view()
                .expect("index query must retain its catalog view");
            if repeated {
                assert!(Arc::ptr_eq(&before.unwrap(), &cached));
            }
        }
    }
    let held = catalog.writer_head_view().unwrap();
    catalog.commit(|tx| tx.drop_index(index)).await.unwrap();
    assert!(held.index_by_id(table, index).is_some());
    assert!(matches!(
        catalog.index_lookup_many(table, index, &[]).await,
        Err(crate::Error::NotFound(_))
    ));
    catalog.close().await.unwrap();
}

#[tokio::test]
async fn a_late_read_view_cannot_replace_a_committed_view() {
    let catalog = open().await;
    let epoch = projection::cache_epoch(catalog.projections());
    let session = catalog.begin_read().await.unwrap();
    let old = catalog.load_head_view(session.handle()).await.unwrap();

    catalog
        .commit(|tx| tx.create_schema("new").map(|_| ()))
        .await
        .unwrap();
    let committed = catalog.writer_head_view().unwrap();
    let late = catalog.head_view(session.handle(), epoch).await.unwrap();
    session.finish();

    assert_eq!(late.current_snapshot().id, old.current_snapshot().id);
    assert!(late.schema_by_name("new").is_none());
    assert!(committed.schema_by_name("new").is_some());
    assert!(Arc::ptr_eq(&committed, &catalog.snapshot().await.unwrap()));
    catalog.close().await.unwrap();
}

#[tokio::test]
async fn recent_row_does_not_fetch_unrequested_live_chunks() {
    let catalog = open().await;
    commit_staged(
        &catalog,
        1,
        "inlined_insert:1",
        vec![
            RowOperation::InlineSchema {
                table_id: 1,
                schema_version: 0,
                arrow_schema: b"schema".to_vec(),
            },
            RowOperation::InlineInsert {
                table_id: 1,
                schema_version: 0,
                begin_snapshot: 1,
                row_id_start: 0,
                row_count: 128,
                arrow_body: b"wanted".to_vec(),
            },
            RowOperation::InlineInsert {
                table_id: 1,
                schema_version: 0,
                begin_snapshot: 1,
                row_id_start: 128,
                row_count: 128,
                arrow_body: b"unrelated".to_vec(),
            },
        ],
    )
    .await;
    assert_eq!(
        catalog
            .recent_row(TableId::new(1), 1)
            .await
            .unwrap()
            .unwrap()
            .offset_in_chunk,
        1
    );
    let tx = catalog.begin_write_tx().await.unwrap();
    tx.delete(
        Key::Inline(InlineKey::Live(InlineOperation::Insert {
            table_id: 1,
            schema_version: 0,
            begin_snapshot: 1,
            chunk_seq: 1,
        }))
        .encode(),
    )
    .unwrap();
    tx.commit().await.unwrap();
    let row = catalog
        .recent_row(TableId::new(1), 1)
        .await
        .unwrap()
        .unwrap();
    assert_eq!(row.chunk_body.as_slice(), b"wanted");
    assert!(
        catalog
            .recent_row(TableId::new(1), 999)
            .await
            .unwrap()
            .is_none()
    );
    catalog.close().await.unwrap();
}

#[tokio::test]
#[allow(clippy::too_many_lines)]
async fn requested_inline_rows_follow_updates_and_wider_chunks() {
    let catalog = open().await;
    commit_staged(
        &catalog,
        1,
        "inlined_insert:1",
        vec![
            RowOperation::InlineSchema {
                table_id: 1,
                schema_version: 0,
                arrow_schema: b"schema".to_vec(),
            },
            RowOperation::InlineInsert {
                table_id: 1,
                schema_version: 0,
                begin_snapshot: 1,
                row_id_start: 10,
                row_count: 2,
                arrow_body: b"old".to_vec(),
            },
        ],
    )
    .await;
    assert_eq!(
        catalog
            .recent_row(TableId::new(1), 10)
            .await
            .unwrap()
            .unwrap()
            .chunk_body
            .as_slice(),
        b"old"
    );
    commit_staged(
        &catalog,
        2,
        "inlined_insert:1,inlined_delete:1",
        vec![
            RowOperation::InlineInlineDelete {
                table_id: 1,
                row_id: 10,
                end_snapshot: 2,
            },
            RowOperation::InlineInlineDelete {
                table_id: 1,
                row_id: 11,
                end_snapshot: 2,
            },
            RowOperation::InlineInsert {
                table_id: 1,
                schema_version: 0,
                begin_snapshot: 2,
                row_id_start: 0,
                row_count: 1024,
                arrow_body: b"replacement".to_vec(),
            },
        ],
    )
    .await;
    let requested = catalog
        .requested_inline_row_ids(TableId::new(1), &[10, 1000, 10, u64::MAX], None)
        .await
        .unwrap();
    assert_eq!(requested, [10, 1000].into_iter().collect());
    let row = catalog
        .recent_row(TableId::new(1), 10)
        .await
        .unwrap()
        .unwrap();
    assert_eq!(row.begin_snapshot.get(), 2);
    assert_eq!(row.chunk_body.as_slice(), b"replacement");
    commit_staged(
        &catalog,
        3,
        "inlined_delete:1",
        vec![RowOperation::InlineInlineDelete {
            table_id: 1,
            row_id: 10,
            end_snapshot: 3,
        }],
    )
    .await;
    assert!(
        catalog
            .recent_row(TableId::new(1), 10)
            .await
            .unwrap()
            .is_none()
    );
    assert!(
        !catalog
            .requested_inline_row_ids(TableId::new(1), &[10], None)
            .await
            .unwrap()
            .contains(&10)
    );
    assert_eq!(
        catalog
            .recent_rows_at(TableId::new(1), crate::SnapshotId::new(1))
            .await
            .unwrap()
            .len(),
        2
    );
    catalog.close().await.unwrap();
}