Repository: MystenLabs/sui
File: crates/sui-pg-db/src/store.rs (Rust)

Code before:
use sui_indexer_alt_framework_store_traits as store;
use sui_sql_macro::sql;

use crate::model::StoredWatermark;
use crate::schema::watermarks;
use crate::{Connection, Db};

pub use sui_indexer_alt_framework_store_traits::Store;

#[async_trait]
impl store::Connection for Connection<'_> {
    async fn init_watermark(
        &mut self,
        pipeline_task: &str,
        first_checkpoint: u64,
    ) -> anyhow::Result<u64> {
        // Create a StoredWatermark directly from CommitterWatermark
        let stored_watermark = StoredWatermark {
            pipeline: pipeline_task.to_string(),
            epoch_hi_inclusive: 0,
            checkpoint_hi_inclusive: first_checkpoint as i64 - 1,
            tx_hi: 0,
            timestamp_ms_hi_inclusive: 0,
            reader_lo: first_checkpoint as i64,
            pruner_timestamp: NaiveDateTime::UNIX_EPOCH,
            pruner_hi: first_checkpoint as i64,
        };

        use diesel::pg::upsert::excluded;
        let checkpoint_hi_inclusive: i64 = diesel::insert_into(watermarks::table)
            .values(&stored_watermark)
            // There is an existing entry, so only write the new `hi` values
            .on_conflict(watermarks::pipeline)
            // Use `do_update` instead of `do_nothing` to return the existing row with `returning`.
            .do_update()
            // When using `do_update`, at least one change needs to be set, so set the pipeline to itself (nothing changes).
            // `excluded` is a virtual table containing the existing row that there was a conflict with.
            .set((watermarks::pipeline.eq(excluded(watermarks::pipeline)),))
            .returning(watermarks::checkpoint_hi_inclusive)
            .get_result(self)
            .await?;

        Ok((checkpoint_hi_inclusive + 1) as u64)
    }

    async fn committer_watermark(
        &mut self,
        pipeline_task: &str,
    ) -> anyhow::Result<Option<store::CommitterWatermark>> {
        let watermark: Option<(i64, i64, i64, i64)> = watermarks::table
            .select((

Review comment (suggestion):
i think as a trait method we should have `init_watermark` accept three checkpoint values for `checkpoint_hi_inclusive`, `reader_lo`, `pruner_hi`, or have the doc comments explain that `init_watermark` is supposed to set these values a particular way

Also, I think we want `pruner_hi <= reader_lo < checkpoint_hi_inclusive`, which should be safe if we handle when `checkpoint_hi_inclusive` is 0

Code after:
use sui_indexer_alt_framework_store_traits as store;
use sui_sql_macro::sql;

use crate::model::StoredWatermark;
use crate::schema::watermarks;
use crate::{Connection, Db};

pub use sui_indexer_alt_framework_store_traits::Store;

#[async_trait]
impl store::Connection for Connection<'_> {
    async fn init_watermark(
        &mut self,
        pipeline_task: &str,
        default_next_checkpoint: u64,
    ) -> anyhow::Result<Option<u64>> {
        let Some(checkpoint_hi_inclusive) = default_next_checkpoint.checked_sub(1) else {
            // Do not create a watermark record with checkpoint_hi_inclusive = -1.
            return Ok(self
                .committer_watermark(pipeline_task)
                .await?
                .map(|w| w.checkpoint_hi_inclusive));
        };

        let stored_watermark = StoredWatermark {
            pipeline: pipeline_task.to_string(),
            epoch_hi_inclusive: 0,
            checkpoint_hi_inclusive: checkpoint_hi_inclusive as i64,
            tx_hi: 0,
            timestamp_ms_hi_inclusive: 0,
            reader_lo: default_next_checkpoint as i64,
            pruner_timestamp: Utc::now().naive_utc(),
            pruner_hi: default_next_checkpoint as i64,
        };

        use diesel::pg::upsert::excluded;
        let checkpoint_hi_inclusive: i64 = diesel::insert_into(watermarks::table)
            .values(&stored_watermark)
            // There is an existing entry, so only write the new `hi` values
            .on_conflict(watermarks::pipeline)
            // Use `do_update` instead of `do_nothing` to return the existing row with `returning`.
            .do_update()
            // When using `do_update`, at least one change needs to be set, so set the pipeline to itself (nothing changes).
            // `excluded` is a virtual table containing the existing row that there was a conflict with.
            .set(watermarks::pipeline.eq(excluded(watermarks::pipeline)))
            .returning(watermarks::checkpoint_hi_inclusive)
            .get_result(self)
            .await?;

        Ok(Some(checkpoint_hi_inclusive as u64))
    }