/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Result;
use async_trait::async_trait;
use bookmarks::Bookmark;
use bookmarks::BookmarkCategory;
use bookmarks::BookmarkKey;
use bookmarks::BookmarkKind;
use bookmarks::BookmarkName;
use bookmarks::BookmarkPagination;
use bookmarks::BookmarkPrefix;
use bookmarks::BookmarkTransaction;
use bookmarks::BookmarkUpdateLog;
use bookmarks::BookmarkUpdateLogEntry;
use bookmarks::BookmarkUpdateReason;
use bookmarks::Bookmarks;
use bookmarks::BookmarksSubscription;
use bookmarks::Freshness;
use cloned::cloned;
use context::CoreContext;
use context::PerfCounterType;
use context::SessionClass;
use futures::future::BoxFuture;
use futures::future::Future;
use futures::future::FutureExt;
use futures::future::TryFutureExt;
use futures::stream;
use futures::stream::BoxStream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use futures_watchdog::WatchdogExt;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use mononoke_types::Timestamp;
use rand::Rng;
use sql::Connection;
use sql_ext::mononoke_queries;
use sql_ext::SqlConnections;
use stats::prelude::*;

use crate::subscription::SqlBookmarksSubscription;
use crate::transaction::SqlBookmarksTransaction;

define_stats! {
    prefix = "mononoke.dbbookmarks";
    list: timeseries(Rate, Sum),
    list_maybe_stale: timeseries(Rate, Sum),
    list_wbc: timeseries(Rate, Sum),
    list_maybe_stale_wbc: timeseries(Rate, Sum),
    get_bookmark: timeseries(Rate, Sum),
}

mononoke_queries! {
    pub(crate) read SelectBookmark(repo_id: RepositoryId, name: BookmarkName, category: BookmarkCategory) -> (ChangesetId, Option<u64>) {
        "SELECT changeset_id, log_id
         FROM bookmarks
         WHERE repo_id = {repo_id}
           AND name = {name}
           AND category = {category}
         LIMIT 1"
    }

    read SelectAll(
        repo_id: RepositoryId,
        limit: u64,
        >list kinds: BookmarkKind
        >list categories: BookmarkCategory
    ) -> (BookmarkName, BookmarkCategory, BookmarkKind, ChangesetId, Option<u64>) {
        "SELECT name, category, hg_kind, changeset_id, log_id
         FROM bookmarks
         WHERE repo_id = {repo_id}
           AND category IN {categories}
           AND hg_kind IN {kinds}
         ORDER BY name ASC
         LIMIT {limit}"
    }

    pub(crate) read SelectAllUnordered(
        repo_id: RepositoryId,
        limit: u64,
        tok: i32,
        >list kinds: BookmarkKind
        >list categories: BookmarkCategory
    ) -> (BookmarkName, BookmarkCategory, BookmarkKind, ChangesetId, Option<u64>, i32) {
        "SELECT name, category, hg_kind, changeset_id, log_id, {tok}
         FROM bookmarks
         WHERE repo_id = {repo_id}
           AND category IN {categories}
           AND hg_kind IN {kinds}
         LIMIT {limit}"
    }

    read SelectAllAfter(
        repo_id: RepositoryId,
        after: BookmarkName,
        limit: u64,
        >list kinds: BookmarkKind
        >list categories: BookmarkCategory
    ) -> (BookmarkName, BookmarkCategory, BookmarkKind, ChangesetId, Option<u64>) {
        "SELECT name, category, hg_kind, changeset_id, log_id
         FROM bookmarks
         WHERE repo_id = {repo_id}
           AND name > {after}
           AND category IN {categories}
           AND hg_kind IN {kinds}
         ORDER BY name ASC
         LIMIT {limit}"
    }

    read SelectByPrefix(
        repo_id: RepositoryId,
        prefix_like_pattern: String,
        escape_character: &str,
        limit: u64,
        >list kinds: BookmarkKind
        >list categories: BookmarkCategory
    ) -> (BookmarkName, BookmarkCategory, BookmarkKind, ChangesetId, Option<u64>) {
        "SELECT name, category, hg_kind, changeset_id, log_id
         FROM bookmarks
         WHERE repo_id = {repo_id}
           AND name LIKE {prefix_like_pattern} ESCAPE {escape_character}
           AND category IN {categories}
           AND hg_kind IN {kinds}
         ORDER BY name ASC
         LIMIT {limit}"
    }

    read SelectByPrefixUnordered(
        repo_id: RepositoryId,
        prefix_like_pattern: String,
        escape_character: &str,
        limit: u64,
        >list kinds: BookmarkKind
        >list categories: BookmarkCategory
    ) -> (BookmarkName, BookmarkCategory, BookmarkKind, ChangesetId, Option<u64>) {
        "SELECT name, category, hg_kind, changeset_id, log_id
         FROM bookmarks
         WHERE repo_id = {repo_id}
           AND name LIKE {prefix_like_pattern} ESCAPE {escape_character}
           AND category IN {categories}
           AND hg_kind IN {kinds}
         LIMIT {limit}"
    }

    read SelectByPrefixAfter(
        repo_id: RepositoryId,
        prefix_like_pattern: String,
        escape_character: &str,
        after: BookmarkName,
        limit: u64,
        >list kinds:BookmarkKind
        >list categories: BookmarkCategory
    ) -> (BookmarkName, BookmarkCategory, BookmarkKind, ChangesetId, Option<u64>) {
        "SELECT name, category, hg_kind, changeset_id, log_id
         FROM bookmarks
         WHERE repo_id = {repo_id}
           AND name LIKE {prefix_like_pattern} ESCAPE {escape_character}
           AND name > {after}
           AND category IN {categories}
           AND hg_kind IN {kinds}
         ORDER BY name ASC
         LIMIT {limit}"
    }

    read SelectAfterLogId(
        repo_id: RepositoryId,
        log_id: u64,
    ) -> (BookmarkName, BookmarkCategory, BookmarkKind, ChangesetId, u64) {
        "SELECT name, category, hg_kind, changeset_id, log_id
         FROM bookmarks
         WHERE repo_id = {repo_id}
           AND log_id IS NOT NULL
           AND log_id > {log_id}"
    }

    read ReadNextBookmarkLogEntries(min_id: u64, repo_id: RepositoryId, limit: u64) -> (
        i64, RepositoryId, BookmarkName, BookmarkCategory, Option<ChangesetId>, Option<ChangesetId>,
        BookmarkUpdateReason, Timestamp
    ) {
        "SELECT id, repo_id, name, category, to_changeset_id, from_changeset_id, reason, timestamp
         FROM bookmarks_update_log
         WHERE id > {min_id} AND repo_id = {repo_id}
         ORDER BY id asc
         LIMIT {limit}"
    }

    read CountFurtherBookmarkLogEntries(min_id: u64, repo_id: RepositoryId) -> (u64) {
        "SELECT COUNT(*)
        FROM bookmarks_update_log
        WHERE id > {min_id} AND repo_id = {repo_id}"
    }

    read CountFurtherBookmarkLogEntriesByReason(min_id: u64, repo_id: RepositoryId) -> (BookmarkUpdateReason, u64) {
        "SELECT reason, COUNT(*)
        FROM bookmarks_update_log
        WHERE id > {min_id} AND repo_id = {repo_id}
        GROUP BY reason"
    }

    read SkipOverBookmarkLogEntriesWithReason(min_id: u64, repo_id: RepositoryId, reason: BookmarkUpdateReason) -> (u64) {
        // We find the first entry that we _don't_ want to skip.
        // Then we find the first entry that we do want to skip and is immediately before this.
        // We don't allow looking back, so if we're going backwards, nothing happens.
        "
        SELECT id
        FROM bookmarks_update_log
        WHERE
            repo_id = {repo_id} AND
            id > {min_id} AND
            reason = {reason} AND
            id < (
                SELECT id
                FROM bookmarks_update_log
                WHERE
                    repo_id = {repo_id} AND
                    id > {min_id} AND
                    NOT reason = {reason}
                ORDER BY id ASC
                LIMIT 1
            )
        ORDER BY id DESC
        LIMIT 1
        "
    }

    read CountFurtherBookmarkLogEntriesWithoutReason(min_id: u64, repo_id: RepositoryId, reason: BookmarkUpdateReason) -> (u64) {
        "SELECT COUNT(*)
        FROM bookmarks_update_log
        WHERE id > {min_id} AND repo_id = {repo_id} AND NOT reason = {reason}"
    }

    read SelectBookmarkLogs(repo_id: RepositoryId, name: BookmarkName, category: BookmarkCategory, max_records: u32, tok: i32) -> (
        u64, Option<ChangesetId>, BookmarkUpdateReason, Timestamp, i32
    ) {
        "SELECT id, to_changeset_id, reason, timestamp, {tok}
         FROM bookmarks_update_log
         WHERE repo_id = {repo_id}
           AND name = {name}
           AND category = {category}
         ORDER BY id DESC
         LIMIT {max_records}"
    }

    read SelectBookmarkLogsWithTsInRange(
        repo_id: RepositoryId,
        name: BookmarkName,
        category: BookmarkCategory,
        max_records: u32,
        min_ts: Timestamp,
        max_ts: Timestamp
    ) -> (
        u64, Option<ChangesetId>, BookmarkUpdateReason, Timestamp
    ) {
        "SELECT id, to_changeset_id, reason, timestamp
         FROM bookmarks_update_log
         WHERE repo_id = {repo_id}
           AND name = {name}
           AND category = {category}
           AND timestamp >= {min_ts}
           AND timestamp <= {max_ts}
         ORDER BY id DESC
         LIMIT {max_records}"
    }

    read SelectBookmarkLogsWithOffset(repo_id: RepositoryId, name: BookmarkName, category: BookmarkCategory, max_records: u32, offset: u32, tok: i32) -> (
        u64, Option<ChangesetId>, BookmarkUpdateReason, Timestamp, i32
    ) {
        "SELECT id, to_changeset_id, reason, timestamp, {tok}
         FROM bookmarks_update_log
         WHERE repo_id = {repo_id}
           AND name = {name}
           AND category = {category}
         ORDER BY id DESC
         LIMIT {max_records}
         OFFSET {offset}"
    }

    pub(crate) read GetLargestLogId(repo_id: RepositoryId) -> (Option<u64>) {
        "SELECT MAX(id)
         FROM bookmarks_update_log
         WHERE repo_id = {repo_id}"
    }
}

#[facet::facet]
#[derive(Clone)]
pub struct SqlBookmarks {
    pub(crate) repo_id: RepositoryId,
    pub(crate) connections: SqlConnections,
}

impl SqlBookmarks {
    pub(crate) fn new(repo_id: RepositoryId, connections: SqlConnections) -> Self {
        Self {
            repo_id,
            connections,
        }
    }

    pub fn connection(&self, ctx: &CoreContext, freshness: Freshness) -> &Connection {
        match freshness {
            Freshness::MaybeStale => {
                ctx.perf_counters()
                    .increment_counter(PerfCounterType::SqlReadsReplica);
                &self.connections.read_connection
            }
            Freshness::MostRecent => {
                ctx.perf_counters()
                    .increment_counter(PerfCounterType::SqlReadsMaster);
                &self.connections.read_master_connection
            }
        }
    }

    pub fn list_raw(
        &self,
        ctx: &CoreContext,
        freshness: Freshness,
        prefix: &BookmarkPrefix,
        categories: &[BookmarkCategory],
        kinds: &[BookmarkKind],
        pagination: &BookmarkPagination,
        limit: u64,
    ) -> impl Future<Output = Result<Vec<(BookmarkKey, BookmarkKind, ChangesetId, Option<u64>)>>>
    {
        let is_wbc = matches!(
            ctx.session().session_class(),
            SessionClass::WarmBookmarksCache
        );

        let conn = match freshness {
            Freshness::MaybeStale => {
                STATS::list_maybe_stale.add_value(1);
                if is_wbc {
                    STATS::list_maybe_stale_wbc.add_value(1);
                }
                ctx.perf_counters()
                    .increment_counter(PerfCounterType::SqlReadsReplica);
                self.connections.read_connection.clone()
            }
            Freshness::MostRecent => {
                STATS::list.add_value(1);
                if is_wbc {
                    STATS::list_wbc.add_value(2);
                }
                ctx.perf_counters()
                    .increment_counter(PerfCounterType::SqlReadsMaster);
                self.connections.read_master_connection.clone()
            }
        };

        cloned!(pagination, prefix, self.repo_id);
        let kinds: Vec<BookmarkKind> = kinds.to_vec();
        let categories: Vec<_> = categories.to_vec();

        async move {
            let rows = if prefix.is_empty() {
                match pagination {
                    BookmarkPagination::FromStart => {
                        // Sorting is only useful for pagination. If the query returns all bookmark
                        // names, then skip the sorting.
                        if limit == std::u64::MAX {
                            let tok: i32 = rand::thread_rng().gen();
                            SelectAllUnordered::query(
                                &conn,
                                &repo_id,
                                &limit,
                                &tok,
                                &kinds,
                                &categories,
                            )
                            .await?
                            .into_iter()
                            .map(|(name, category, kind, cs_id, log_id, _)| {
                                (
                                    BookmarkKey::with_name_and_category(name, category),
                                    kind,
                                    cs_id,
                                    log_id,
                                )
                            })
                            .collect()
                        } else {
                            SelectAll::query(&conn, &repo_id, &limit, &kinds, &categories)
                                .await?
                                .into_iter()
                                .map(|(name, category, kind, cs_id, log_id)| {
                                    (
                                        BookmarkKey::with_name_and_category(name, category),
                                        kind,
                                        cs_id,
                                        log_id,
                                    )
                                })
                                .collect()
                        }
                    }
                    BookmarkPagination::After(after) => {
                        SelectAllAfter::query(&conn, &repo_id, &after, &limit, &kinds, &categories)
                            .await?
                            .into_iter()
                            .map(|(name, category, kind, cs_id, log_id)| {
                                (
                                    BookmarkKey::with_name_and_category(name, category),
                                    kind,
                                    cs_id,
                                    log_id,
                                )
                            })
                            .collect()
                    }
                }
            } else {
                let prefix_like_pattern = prefix.to_escaped_sql_like_pattern();
                match pagination {
                    BookmarkPagination::FromStart => {
                        if limit == std::u64::MAX {
                            SelectByPrefixUnordered::query(
                                &conn,
                                &repo_id,
                                &prefix_like_pattern,
                                &"\\",
                                &limit,
                                &kinds,
                                &categories,
                            )
                            .await?
                            .into_iter()
                            .map(|(name, category, kind, cs_id, log_id)| {
                                (
                                    BookmarkKey::with_name_and_category(name, category),
                                    kind,
                                    cs_id,
                                    log_id,
                                )
                            })
                            .collect()
                        } else {
                            SelectByPrefix::query(
                                &conn,
                                &repo_id,
                                &prefix_like_pattern,
                                &"\\",
                                &limit,
                                &kinds,
                                &categories,
                            )
                            .await?
                            .into_iter()
                            .map(|(name, category, kind, cs_id, log_id)| {
                                (
                                    BookmarkKey::with_name_and_category(name, category),
                                    kind,
                                    cs_id,
                                    log_id,
                                )
                            })
                            .collect()
                        }
                    }
                    BookmarkPagination::After(after) => SelectByPrefixAfter::query(
                        &conn,
                        &repo_id,
                        &prefix_like_pattern,
                        &"\\",
                        &after,
                        &limit,
                        &kinds,
                        &categories,
                    )
                    .await?
                    .into_iter()
                    .map(|(name, category, kind, cs_id, log_id)| {
                        (
                            BookmarkKey::with_name_and_category(name, category),
                            kind,
                            cs_id,
                            log_id,
                        )
                    })
                    .collect(),
                }
            };

            Ok(rows)
        }
    }

    pub fn get_raw(
        &self,
        ctx: CoreContext,
        key: &BookmarkKey,
    ) -> impl Future<Output = Result<Option<(ChangesetId, Option<u64>)>>> + 'static {
        STATS::get_bookmark.add_value(1);
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsMaster);
        let conn = self.connections.read_master_connection.clone();
        cloned!(self.repo_id, key);
        async move {
            let rows = SelectBookmark::query(&conn, &repo_id, key.name(), key.category()).await?;
            Ok(rows.into_iter().next())
        }
    }
}

#[async_trait]
impl Bookmarks for SqlBookmarks {
    fn list(
        &self,
        ctx: CoreContext,
        freshness: Freshness,
        prefix: &BookmarkPrefix,
        categories: &[BookmarkCategory],
        kinds: &[BookmarkKind],
        pagination: &BookmarkPagination,
        limit: u64,
    ) -> BoxStream<'static, Result<(Bookmark, ChangesetId)>> {
        let fut = self.list_raw(
            &ctx, freshness, prefix, categories, kinds, pagination, limit,
        );

        async move {
            let rows = fut.await?;

            Ok(stream::iter(rows.into_iter().map(|row| {
                let (key, kind, changeset_id, _log_id) = row;
                Ok((Bookmark::new(key, kind), changeset_id))
            })))
        }
        .try_flatten_stream()
        .boxed()
    }

    fn get(
        &self,
        ctx: CoreContext,
        name: &BookmarkKey,
    ) -> BoxFuture<'static, Result<Option<ChangesetId>>> {
        self.get_raw(ctx, name)
            .map_ok(|maybe_row| maybe_row.map(|(cs_id, _log_id)| cs_id))
            .boxed()
    }

    async fn create_subscription(
        &self,
        ctx: &CoreContext,
        freshness: Freshness,
    ) -> Result<Box<dyn BookmarksSubscription>> {
        let sub = SqlBookmarksSubscription::create(ctx, self.clone(), freshness)
            .await
            .context("Failed to create SqlBookmarksSubscription")?;

        Ok(Box::new(sub))
    }

    fn create_transaction(&self, ctx: CoreContext) -> Box<dyn BookmarkTransaction> {
        Box::new(SqlBookmarksTransaction::new(
            ctx,
            self.connections.write_connection.clone(),
            self.repo_id.clone(),
        ))
    }
}

impl BookmarkUpdateLog for SqlBookmarks {
    fn list_bookmark_log_entries(
        &self,
        ctx: CoreContext,
        key: BookmarkKey,
        max_rec: u32,
        offset: Option<u32>,
        freshness: Freshness,
    ) -> BoxStream<'static, Result<(u64, Option<ChangesetId>, BookmarkUpdateReason, Timestamp)>>
    {
        let conn = if freshness == Freshness::MostRecent {
            ctx.perf_counters()
                .increment_counter(PerfCounterType::SqlReadsMaster);
            self.connections.read_master_connection.clone()
        } else {
            ctx.perf_counters()
                .increment_counter(PerfCounterType::SqlReadsReplica);
            self.connections.read_connection.clone()
        };
        let repo_id = self.repo_id;

        async move {
            let tok: i32 = rand::thread_rng().gen();

            let rows = match offset {
                Some(offset) => {
                    SelectBookmarkLogsWithOffset::query(
                        &conn,
                        &repo_id,
                        key.name(),
                        key.category(),
                        &max_rec,
                        &offset,
                        &tok,
                    )
                    .await?
                }
                None => {
                    SelectBookmarkLogs::query(
                        &conn,
                        &repo_id,
                        key.name(),
                        key.category(),
                        &max_rec,
                        &tok,
                    )
                    .await?
                }
            };
            Ok(stream::iter(
                rows.into_iter()
                    .map(|(from_id, to_id, reason, ts, _)| (from_id, to_id, reason, ts))
                    .map(Ok),
            ))
        }
        .try_flatten_stream()
        .boxed()
    }

    fn list_bookmark_log_entries_ts_in_range(
        &self,
        ctx: CoreContext,
        key: BookmarkKey,
        max_rec: u32,
        min_ts: Timestamp,
        max_ts: Timestamp,
    ) -> BoxStream<'static, Result<(u64, Option<ChangesetId>, BookmarkUpdateReason, Timestamp)>>
    {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);

        let conn = self.connections.read_connection.clone();
        let repo_id = self.repo_id;

        async move {
            let rows = SelectBookmarkLogsWithTsInRange::query(
                &conn,
                &repo_id,
                key.name(),
                key.category(),
                &max_rec,
                &min_ts,
                &max_ts,
            )
            .await?;
            Ok(stream::iter(rows.into_iter().map(Ok)))
        }
        .try_flatten_stream()
        .boxed()
    }

    fn count_further_bookmark_log_entries(
        &self,
        ctx: CoreContext,
        id: u64,
        maybe_exclude_reason: Option<BookmarkUpdateReason>,
    ) -> BoxFuture<'static, Result<u64>> {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);
        let conn = self.connections.read_connection.clone();
        let repo_id = self.repo_id;

        async move {
            let entries = match maybe_exclude_reason {
                Some(ref r) => {
                    CountFurtherBookmarkLogEntriesWithoutReason::query(&conn, &id, &repo_id, r)
                        .await?
                }
                None => CountFurtherBookmarkLogEntries::query(&conn, &id, &repo_id).await?,
            };
            match entries.into_iter().next() {
                Some(count) => Ok(count.0),
                None => {
                    let extra = match maybe_exclude_reason {
                        Some(..) => "without reason",
                        None => "",
                    };
                    Err(anyhow!(
                        "Failed to query further bookmark log entries{}",
                        extra
                    ))
                }
            }
        }
        .boxed()
    }

    fn count_further_bookmark_log_entries_by_reason(
        &self,
        ctx: CoreContext,
        id: u64,
    ) -> BoxFuture<'static, Result<Vec<(BookmarkUpdateReason, u64)>>> {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);
        let conn = self.connections.read_connection.clone();
        let repo_id = self.repo_id;
        async move {
            let entries =
                CountFurtherBookmarkLogEntriesByReason::query(&conn, &id, &repo_id).await?;
            Ok(entries.into_iter().collect())
        }
        .boxed()
    }

    fn skip_over_bookmark_log_entries_with_reason(
        &self,
        ctx: CoreContext,
        id: u64,
        reason: BookmarkUpdateReason,
    ) -> BoxFuture<'static, Result<Option<u64>>> {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);
        let conn = self.connections.read_connection.clone();
        cloned!(self.repo_id, reason);
        async move {
            let entries =
                SkipOverBookmarkLogEntriesWithReason::query(&conn, &id, &repo_id, &reason).await?;
            Ok(entries.first().map(|entry| entry.0))
        }
        .boxed()
    }

    fn read_next_bookmark_log_entries_same_bookmark_and_reason(
        &self,
        ctx: CoreContext,
        id: u64,
        limit: u64,
    ) -> BoxStream<'static, Result<BookmarkUpdateLogEntry>> {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);
        let conn = self.connections.read_connection.clone();
        let repo_id = self.repo_id;

        async move {
            let entries = ReadNextBookmarkLogEntries::query(&conn, &id, &repo_id, &limit)
                .watched(ctx.logger())
                .await?;

            let homogenous_entries: Vec<_> = match entries.first().cloned() {
                Some(first_entry) => {
                    // Note: types are explicit here to protect us from query behavior change
                    //       when tuple items 2, 3, or 6 become something else, and we still succeed
                    //       compiling everything because of the type inference
                    let first_name: &BookmarkKey = &BookmarkKey::with_name_and_category(
                        first_entry.2.clone(),
                        first_entry.3.clone(),
                    );
                    let first_reason: &BookmarkUpdateReason = &first_entry.6;
                    entries
                        .into_iter()
                        .take_while(|entry| {
                            let name: &BookmarkKey = &BookmarkKey::with_name_and_category(
                                entry.2.clone(),
                                entry.3.clone(),
                            );
                            let reason: &BookmarkUpdateReason = &entry.6;
                            name == first_name && reason == first_reason
                        })
                        .collect()
                }
                None => entries.into_iter().collect(),
            };
            Ok(
                stream::iter(homogenous_entries.into_iter().map(Ok)).and_then(|entry| async move {
                    let (id, repo_id, name, category, to_cs_id, from_cs_id, reason, timestamp) =
                        entry;
                    Ok(BookmarkUpdateLogEntry {
                        id,
                        repo_id,
                        bookmark_name: BookmarkKey::with_name_and_category(name, category),
                        to_changeset_id: to_cs_id,
                        from_changeset_id: from_cs_id,
                        reason,
                        timestamp,
                    })
                }),
            )
        }
        .try_flatten_stream()
        .boxed()
    }

    fn read_next_bookmark_log_entries(
        &self,
        ctx: CoreContext,
        id: u64,
        limit: u64,
        freshness: Freshness,
    ) -> BoxStream<'static, Result<BookmarkUpdateLogEntry>> {
        let connection = if freshness == Freshness::MostRecent {
            ctx.perf_counters()
                .increment_counter(PerfCounterType::SqlReadsMaster);
            self.connections.read_master_connection.clone()
        } else {
            ctx.perf_counters()
                .increment_counter(PerfCounterType::SqlReadsReplica);
            self.connections.read_connection.clone()
        };

        let repo_id = self.repo_id;

        async move {
            let entries =
                ReadNextBookmarkLogEntries::query(&connection, &id, &repo_id, &limit).await?;

            Ok(
                stream::iter(entries.into_iter().map(Ok)).and_then(|entry| async move {
                    let (id, repo_id, name, category, to_cs_id, from_cs_id, reason, timestamp) =
                        entry;
                    Ok(BookmarkUpdateLogEntry {
                        id,
                        repo_id,
                        bookmark_name: BookmarkKey::with_name_and_category(name, category),
                        to_changeset_id: to_cs_id,
                        from_changeset_id: from_cs_id,
                        reason,
                        timestamp,
                    })
                }),
            )
        }
        .try_flatten_stream()
        .boxed()
    }

    fn get_largest_log_id(
        &self,
        ctx: CoreContext,
        freshness: Freshness,
    ) -> BoxFuture<'static, Result<Option<u64>>> {
        let connection = if freshness == Freshness::MostRecent {
            ctx.perf_counters()
                .increment_counter(PerfCounterType::SqlReadsMaster);
            self.connections.read_master_connection.clone()
        } else {
            ctx.perf_counters()
                .increment_counter(PerfCounterType::SqlReadsReplica);
            self.connections.read_connection.clone()
        };
        let repo_id = self.repo_id;

        async move {
            let entries = GetLargestLogId::query(&connection, &repo_id).await?;
            let entry = entries.into_iter().next();
            match entry {
                Some(count) => Ok(count.0),
                None => Err(anyhow!("Failed to query largest log id")),
            }
        }
        .boxed()
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashMap;
use std::time::Duration;
use std::time::Instant;

use anyhow::bail;
use anyhow::Context;
use anyhow::Result;
use async_trait::async_trait;
use bookmarks::BookmarkCategory;
use bookmarks::BookmarkKey;
use bookmarks::BookmarkKind;
use bookmarks::BookmarkName;
use bookmarks::BookmarksSubscription;
use bookmarks::Freshness;
use context::CoreContext;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use rand::Rng;
use slog::warn;
use sql_ext::mononoke_queries;
use stats::prelude::*;
use tunables::tunables;

use crate::store::GetLargestLogId;
use crate::store::SelectAllUnordered;
use crate::store::SqlBookmarks;

define_stats! {
    prefix = "mononoke.dbbookmarks.subscription";
    aged_out: timeseries(Sum),
    master_protected: timeseries(Sum),
}

const DEFAULT_SUBSCRIPTION_MAX_AGE_MS: u64 = 600_000; // 10 minutes.

#[derive(Clone)]
pub struct SqlBookmarksSubscription {
    sql_bookmarks: SqlBookmarks,
    freshness: Freshness,
    log_id: u64,
    bookmarks: HashMap<BookmarkKey, (ChangesetId, BookmarkKind)>,
    last_refresh: Instant,
}

impl SqlBookmarksSubscription {
    pub async fn create(
        ctx: &CoreContext,
        sql_bookmarks: SqlBookmarks,
        freshness: Freshness,
    ) -> Result<Self> {
        // We assume that we're using the default transaction isolation levels (which is repeatable
        // read). This means that the first select here will establish a snapshot, and we'll see
        // bookmarks updated prior to this log id in the second query.

        let conn = sql_bookmarks.connection(ctx, freshness);
        let txn = conn
            .start_transaction()
            .await
            .context("Failed to start bookmarks read transaction")?;

        let (txn, log_id_rows) =
            GetLargestLogId::query_with_transaction(txn, &sql_bookmarks.repo_id)
                .await
                .context("Failed to read log id")?;

        // Our ids start at 1 so we can default log_id to zero if it's missing.
        let log_id = log_id_rows
            .into_iter()
            .map(|r| r.0)
            .next()
            .flatten()
            .unwrap_or(0);

        let tok: i32 = rand::thread_rng().gen();
        let (txn, bookmarks) = SelectAllUnordered::query_with_transaction(
            txn,
            &sql_bookmarks.repo_id,
            &std::u64::MAX,
            &tok,
            BookmarkKind::ALL_PUBLISHING,
            BookmarkCategory::ALL,
        )
        .await
        .context("Failed to query bookmarks")?;

        // Cleanly close this transaction. No need to commit: we're not making any changes here.
        txn.rollback()
            .await
            .context("Failed to close transaction")?;

        let bookmarks = bookmarks
            .into_iter()
            .map(|(name, category, kind, cs_id, _log_id, _tok)| {
                (
                    BookmarkKey::with_name_and_category(name, category),
                    (cs_id, kind),
                )
            })
            .collect();

        Ok(Self {
            sql_bookmarks,
            freshness,
            log_id,
            bookmarks,
            last_refresh: Instant::now(),
        })
    }

    fn has_aged_out(&self) -> bool {
        let max_age_ms = match tunables()
            .bookmark_subscription_max_age_ms()
            .unwrap_or_default()
            .try_into()
        {
            Ok(duration) if duration > 0 => duration,
            _ => DEFAULT_SUBSCRIPTION_MAX_AGE_MS,
        };

        self.last_refresh.elapsed() > Duration::from_millis(max_age_ms)
    }
}

#[async_trait]
impl BookmarksSubscription for SqlBookmarksSubscription {
    async fn refresh(&mut self, ctx: &CoreContext) -> Result<()> {
        if self.has_aged_out() {
            warn!(
                ctx.logger(),
                "BookmarksSubscription has aged out! Last refresh was {:?} ago.",
                self.last_refresh.elapsed()
            );

            STATS::aged_out.add_value(1);

            *self = Self::create(ctx, self.sql_bookmarks.clone(), self.freshness)
                .await
                .context("Failed to re-initialize expired BookmarksSubscription")?;
            return Ok(());
        }

        let conn = self.sql_bookmarks.connection(ctx, self.freshness);

        let changes =
            SelectUpdatedBookmarks::query(conn, &self.sql_bookmarks.repo_id, &self.log_id)
                .await
                .with_context(|| {
                    format!("Failed to select updated bookmarks after {}", self.log_id)
                })?;

        let mut max_log_id = None;
        let mut updates = HashMap::new();

        for (log_id, name, category, kind, cs_id) in changes {
            let bookmark = BookmarkKey::with_name_and_category(name, category);
            // kind & cs_id come from the same table (bookmarks) and they're not nullable there, so
            // if one is missing, that means the join didn't find anything, and the one must be
            // missing too.
            let value = match (kind, cs_id) {
                (Some(kind), Some(cs_id)) => Some((cs_id, kind)),
                (None, None) => None,
                _ => bail!("Inconsistent data for bookmark: {}", bookmark),
            };

            max_log_id = std::cmp::max(max_log_id, Some(log_id));

            // NOTE: We get the updates in DESC-ending order, so we'll always find the curent
            // bookmark state first.
            updates.entry(bookmark).or_insert(value);
        }

        for (book, maybe_value) in updates.into_iter() {
            match maybe_value {
                Some(value) => {
                    self.bookmarks.insert(book, value);
                }
                None => {
                    if tunables()
                        .bookmark_subscription_protect_master()
                        .unwrap_or_default()
                    {
                        if book.as_str() == "master" {
                            warn!(
                                ctx.logger(),
                                "BookmarksSubscription: protect master kicked in!"
                            );
                            STATS::master_protected.add_value(1);
                            continue;
                        }
                    }

                    self.bookmarks.remove(&book);
                }
            }
        }

        if let Some(max_log_id) = max_log_id {
            self.log_id = max_log_id;
        }

        self.last_refresh = Instant::now();

        Ok(())
    }

    fn bookmarks(&self) -> &HashMap<BookmarkKey, (ChangesetId, BookmarkKind)> {
        &self.bookmarks
    }
}

mononoke_queries! {
    read SelectUpdatedBookmarks(
        repo_id: RepositoryId,
        log_id: u64
    ) -> (u64, BookmarkName, BookmarkCategory, Option<BookmarkKind>, Option<ChangesetId>) {
        mysql("
        SELECT bookmarks_update_log.id, bookmarks_update_log.name, bookmarks_update_log.category, bookmarks.hg_kind, bookmarks.changeset_id
        FROM bookmarks_update_log
        LEFT JOIN bookmarks
            FORCE INDEX (repo_id_log_id)
            ON  bookmarks.repo_id = bookmarks_update_log.repo_id
            AND bookmarks.log_id  = bookmarks_update_log.id
        WHERE bookmarks_update_log.repo_id = {repo_id} AND bookmarks_update_log.id > {log_id}
        ORDER BY bookmarks_update_log.id DESC
        ")
        sqlite("
        SELECT bookmarks_update_log.id, bookmarks_update_log.name, bookmarks_update_log.category, bookmarks.hg_kind, bookmarks.changeset_id
        FROM bookmarks_update_log
        LEFT JOIN bookmarks
            ON  bookmarks.repo_id = bookmarks_update_log.repo_id
            AND bookmarks.log_id  = bookmarks_update_log.id
        WHERE bookmarks_update_log.repo_id = {repo_id} AND bookmarks_update_log.id > {log_id}
        ORDER BY bookmarks_update_log.id DESC
        ")
    }
}

#[cfg(test)]
mod test {
    use bookmarks::BookmarkUpdateReason;
    use bookmarks::Bookmarks;
    use fbinit::FacebookInit;
    use maplit::hashmap;
    use mononoke_types_mocks::changesetid::ONES_CSID;
    use mononoke_types_mocks::repo::REPO_ZERO;
    use sql_construct::SqlConstruct;
    use tunables::MononokeTunables;

    use super::*;
    use crate::builder::SqlBookmarksBuilder;

    #[fbinit::test]
    async fn test_age_out(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let bookmarks = SqlBookmarksBuilder::with_sqlite_in_memory()?.with_repo_id(REPO_ZERO);
        let conn = bookmarks.connections.write_connection.clone();

        let mut sub =
            SqlBookmarksSubscription::create(&ctx, bookmarks.clone(), Freshness::MostRecent)
                .await?;

        // Insert a bookmark, but without going throguh the log. This won't happen in prod.
        // However, for the purposes of this test, it makes the update invisible to the
        // subscription, and only a full refresh will find it.
        let book = BookmarkKey::new("book")?;
        let rows = vec![(
            &bookmarks.repo_id,
            &book,
            &ONES_CSID,
            &BookmarkKind::Publishing,
        )];
        crate::transaction::insert_bookmarks(&conn, rows).await?;

        sub.refresh(&ctx).await?;
        assert_eq!(*sub.bookmarks(), hashmap! {});

        sub.last_refresh = Instant::now()
            .checked_sub(Duration::from_millis(DEFAULT_SUBSCRIPTION_MAX_AGE_MS + 1))
            .ok_or(anyhow::anyhow!("Invalid duration subtraction"))?;

        sub.refresh(&ctx).await?;
        assert_eq!(
            *sub.bookmarks(),
            hashmap! { book => (ONES_CSID, BookmarkKind::Publishing)}
        );

        Ok(())
    }

    #[fbinit::test]
    async fn test_protect_master(fb: FacebookInit) -> Result<()> {
        let protect_master = MononokeTunables::default();
        protect_master
            .update_bools(&hashmap! {"bookmark_subscription_protect_master".to_string() => true});

        let ctx = CoreContext::test_mock(fb);
        let bookmarks = SqlBookmarksBuilder::with_sqlite_in_memory()?.with_repo_id(REPO_ZERO);
        let conn = bookmarks.connections.write_connection.clone();

        let book = BookmarkKey::new("master")?;

        let rows = vec![(
            &bookmarks.repo_id,
            &book,
            &ONES_CSID,
            &BookmarkKind::Publishing,
        )];
        crate::transaction::insert_bookmarks(&conn, rows).await?;

        let mut sub1 =
            SqlBookmarksSubscription::create(&ctx, bookmarks.clone(), Freshness::MostRecent)
                .await?;

        let mut sub2 =
            SqlBookmarksSubscription::create(&ctx, bookmarks.clone(), Freshness::MostRecent)
                .await?;

        let mut txn = bookmarks.create_transaction(ctx.clone());
        txn.force_delete(&book, BookmarkUpdateReason::TestMove)?;
        txn.commit().await?;

        tunables::with_tunables_async(protect_master, sub1.refresh(&ctx)).await?;
        assert!(sub1.bookmarks.get(&book).is_some());

        sub2.refresh(&ctx).await?;
        assert!(sub2.bookmarks.get(&book).is_none());

        Ok(())
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashSet;
use std::sync::Arc;

use anyhow::anyhow;
use anyhow::Result;
use bookmarks::BookmarkCategory;
use bookmarks::BookmarkKey;
use bookmarks::BookmarkKind;
use bookmarks::BookmarkName;
use bookmarks::BookmarkTransaction;
use bookmarks::BookmarkTransactionError;
use bookmarks::BookmarkTransactionHook;
use bookmarks::BookmarkUpdateReason;
use context::CoreContext;
use context::PerfCounterType;
use futures::future;
use futures::future::BoxFuture;
use futures::future::FutureExt;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use mononoke_types::Timestamp;
use sql::Connection;
use sql::Transaction as SqlTransaction;
use sql_ext::mononoke_queries;
use stats::prelude::*;

use crate::store::SelectBookmark;

const MAX_BOOKMARK_TRANSACTION_ATTEMPT_COUNT: usize = 5;

define_stats! {
    prefix = "mononoke.dbbookmarks";
    bookmarks_update_log_insert_success: timeseries(Rate, Sum),
    bookmarks_update_log_insert_success_attempt_count: timeseries(Rate, Average, Sum),
    bookmarks_update_log_insert_retry: timeseries(Rate, Sum),
    bookmarks_insert_retryable_error: timeseries(Rate, Sum),
    bookmarks_insert_retryable_error_attempt_count: timeseries(Rate, Average, Sum),
    bookmarks_insert_logic_error: timeseries(Rate, Sum),
    bookmarks_insert_logic_error_attempt_count: timeseries(Rate, Average, Sum),
    bookmarks_insert_other_error: timeseries(Rate, Sum),
    bookmarks_insert_other_error_attempt_count: timeseries(Rate, Average, Sum),
}

mononoke_queries! {
    write ReplaceBookmarks(
        values: (repo_id: RepositoryId, log_id: Option<u64>, name: BookmarkName, category: BookmarkCategory, changeset_id: ChangesetId)
    ) {
        none,
        "REPLACE INTO bookmarks (repo_id, log_id, name, category, changeset_id) VALUES {values}"
    }

    write InsertBookmarks(
        values: (repo_id: RepositoryId, log_id: Option<u64>, name: BookmarkName, category: BookmarkCategory, changeset_id: ChangesetId, kind: BookmarkKind)
    ) {
        insert_or_ignore,
        "{insert_or_ignore} INTO bookmarks (repo_id, log_id, name, category, changeset_id, hg_kind) VALUES {values}"
    }

    write UpdateBookmark(
        repo_id: RepositoryId,
        log_id: Option<u64>,
        name: BookmarkName,
        category: BookmarkCategory,
        old_id: ChangesetId,
        new_id: ChangesetId,
        >list kinds: BookmarkKind
    ) {
        none,
        "UPDATE bookmarks
         SET log_id = {log_id}, changeset_id = {new_id}
         WHERE repo_id = {repo_id}
           AND name = {name}
           AND category = {category}
           AND changeset_id = {old_id}
           AND hg_kind IN {kinds}"
    }

    write DeleteBookmark(repo_id: RepositoryId, name: BookmarkName, category: BookmarkCategory) {
        none,
        "DELETE FROM bookmarks
         WHERE repo_id = {repo_id}
           AND name = {name}
           AND category = {category}"
    }

    write DeleteBookmarkIf(repo_id: RepositoryId, name: BookmarkName, category: BookmarkCategory, changeset_id: ChangesetId) {
        none,
        "DELETE FROM bookmarks
         WHERE repo_id = {repo_id}
           AND name = {name}
           AND category = {category}
           AND changeset_id = {changeset_id}"
    }

    read FindMaxBookmarkLogId(repo_id: RepositoryId) -> (Option<u64>) {
        "SELECT MAX(id) FROM bookmarks_update_log WHERE repo_id = {repo_id}"
    }

    write AddBookmarkLog(
        values: (
            id: u64,
            repo_id: RepositoryId,
            name: BookmarkName,
            category: BookmarkCategory,
            from_changeset_id: Option<ChangesetId>,
            to_changeset_id: Option<ChangesetId>,
            reason: BookmarkUpdateReason,
            timestamp: Timestamp,
        ),
    ) {
        none,
        "INSERT INTO bookmarks_update_log
         (id, repo_id, name, category, from_changeset_id, to_changeset_id, reason, timestamp)
         VALUES {values}"
    }
}

struct NewUpdateLogEntry {
    /// The old bookmarked changeset (if known)
    old: Option<ChangesetId>,

    /// The new bookmarked changeset (or None if the bookmark is being
    /// deleted).
    new: Option<ChangesetId>,

    /// The reason for the update.
    reason: BookmarkUpdateReason,
}

impl NewUpdateLogEntry {
    fn new(
        old: Option<ChangesetId>,
        new: Option<ChangesetId>,
        reason: BookmarkUpdateReason,
    ) -> Result<NewUpdateLogEntry> {
        Ok(NewUpdateLogEntry { old, new, reason })
    }
}

struct SqlBookmarksTransactionPayload {
    /// The repository we are updating.
    repo_id: RepositoryId,

    /// Operations to force-set a bookmark to a changeset.
    force_sets: Vec<(BookmarkKey, ChangesetId, NewUpdateLogEntry)>,

    /// Operations to create a bookmark.
    creates: Vec<(
        BookmarkKey,
        ChangesetId,
        BookmarkKind,
        Option<NewUpdateLogEntry>,
    )>,

    /// Operations to update a bookmark from an old id to a new id, provided
    /// it has a matching kind.
    updates: Vec<(
        BookmarkKey,
        ChangesetId,
        ChangesetId,
        &'static [BookmarkKind],
        Option<NewUpdateLogEntry>,
    )>,

    /// Operations to force-delete a bookmark.
    force_deletes: Vec<(BookmarkKey, NewUpdateLogEntry)>,

    /// Operations to delete a bookmark with an old id.
    deletes: Vec<(BookmarkKey, ChangesetId, Option<NewUpdateLogEntry>)>,
}

/// Structure representing the log entries to insert when executing a
/// SqlBookmarksTransactionPayload.
struct TransactionLogUpdates<'a> {
    next_log_id: u64,
    log_entries: Vec<(u64, &'a BookmarkKey, &'a NewUpdateLogEntry)>,
}

impl<'a> TransactionLogUpdates<'a> {
    fn new(next_log_id: u64) -> Self {
        Self {
            next_log_id,
            log_entries: Vec::new(),
        }
    }

    fn push_log_entry(&mut self, bookmark: &'a BookmarkKey, entry: &'a NewUpdateLogEntry) -> u64 {
        let id = self.next_log_id;
        self.log_entries.push((id, bookmark, entry));
        self.next_log_id += 1;
        id
    }
}

impl SqlBookmarksTransactionPayload {
    fn new(repo_id: RepositoryId) -> Self {
        SqlBookmarksTransactionPayload {
            repo_id,
            force_sets: Vec::new(),
            creates: Vec::new(),
            updates: Vec::new(),
            force_deletes: Vec::new(),
            deletes: Vec::new(),
        }
    }

    async fn find_next_update_log_id(
        txn: SqlTransaction,
        repo_id: RepositoryId,
    ) -> Result<(SqlTransaction, u64)> {
        let (txn, max_id_entries) =
            FindMaxBookmarkLogId::query_with_transaction(txn, &repo_id).await?;

        let next_id = match &max_id_entries[..] {
            [(None,)] => 1,
            [(Some(max_existing),)] => *max_existing + 1,
            _ => {
                return Err(anyhow!(
                    "FindMaxBookmarkLogId returned multiple entries: {:?}",
                    max_id_entries
                ));
            }
        };
        Ok((txn, next_id))
    }

    async fn store_log<'a>(
        &'a self,
        mut txn: SqlTransaction,
        log: &'a TransactionLogUpdates<'a>,
    ) -> Result<SqlTransaction> {
        let timestamp = Timestamp::now();

        for (id, bookmark, log_entry) in log.log_entries.iter() {
            let data = [(
                id,
                &self.repo_id,
                bookmark.name(),
                bookmark.category(),
                &log_entry.old,
                &log_entry.new,
                &log_entry.reason,
                &timestamp,
            )];
            txn = AddBookmarkLog::query_with_transaction(txn, &data[..])
                .await?
                .0;
        }
        Ok(txn)
    }

    async fn store_force_sets<'op, 'log: 'op>(
        &'log self,
        txn: SqlTransaction,
        log: &'op mut TransactionLogUpdates<'log>,
    ) -> Result<SqlTransaction, BookmarkTransactionError> {
        let mut data = Vec::new();
        for (bookmark, cs_id, log_entry) in self.force_sets.iter() {
            let log_id = log.push_log_entry(bookmark, log_entry);
            data.push((self.repo_id, Some(log_id), bookmark, cs_id));
        }
        let data = data
            .iter()
            .map(|(repo_id, log_id, bookmark, cs_id)| {
                (
                    repo_id,
                    log_id,
                    bookmark.name(),
                    bookmark.category(),
                    *cs_id,
                )
            })
            .collect::<Vec<_>>();
        let (txn, _) = ReplaceBookmarks::query_with_transaction(txn, data.as_slice()).await?;
        Ok(txn)
    }

    async fn store_creates<'op, 'log: 'op>(
        &'log self,
        txn: SqlTransaction,
        log: &'op mut TransactionLogUpdates<'log>,
    ) -> Result<SqlTransaction, BookmarkTransactionError> {
        let mut data = Vec::new();
        for (bookmark, cs_id, kind, maybe_log_entry) in self.creates.iter() {
            let log_id = maybe_log_entry
                .as_ref()
                .map(|log_entry| log.push_log_entry(bookmark, log_entry));
            data.push((self.repo_id, log_id, bookmark, cs_id, kind))
        }
        let data = data
            .iter()
            .map(|(repo_id, log_id, bookmark, cs_id, kind)| {
                (
                    repo_id,
                    log_id,
                    bookmark.name(),
                    bookmark.category(),
                    *cs_id,
                    *kind,
                )
            })
            .collect::<Vec<_>>();
        let rows_to_insert = data.len() as u64;
        let (txn, result) = InsertBookmarks::query_with_transaction(txn, data.as_slice()).await?;
        if result.affected_rows() != rows_to_insert {
            return Err(BookmarkTransactionError::LogicError);
        }
        Ok(txn)
    }

    async fn store_updates<'op, 'log: 'op>(
        &'log self,
        mut txn: SqlTransaction,
        log: &'op mut TransactionLogUpdates<'log>,
    ) -> Result<SqlTransaction, BookmarkTransactionError> {
        for (bookmark, old_cs_id, new_cs_id, kinds, maybe_log_entry) in self.updates.iter() {
            let log_id = maybe_log_entry
                .as_ref()
                .map(|log_entry| log.push_log_entry(bookmark, log_entry));

            if new_cs_id == old_cs_id && log_id.is_none() {
                // This is a no-op update.  Check if the bookmark already points to the correct
                // commit.  If it doesn't, abort the transaction. We need to make this a select
                // query instead of an update, since affected_rows() woud otherwise return 0.
                let (txn_, result) = SelectBookmark::query_with_transaction(
                    txn,
                    &self.repo_id,
                    bookmark.name(),
                    bookmark.category(),
                )
                .await?;
                txn = txn_;
                if result.get(0).map(|row| row.0).as_ref() != Some(new_cs_id) {
                    return Err(BookmarkTransactionError::LogicError);
                }
            } else {
                let (txn_, result) = UpdateBookmark::query_with_transaction(
                    txn,
                    &self.repo_id,
                    &log_id,
                    bookmark.name(),
                    bookmark.category(),
                    old_cs_id,
                    new_cs_id,
                    kinds,
                )
                .await?;
                txn = txn_;
                if result.affected_rows() != 1 {
                    return Err(BookmarkTransactionError::LogicError);
                }
            }
        }
        Ok(txn)
    }

    async fn store_force_deletes<'op, 'log: 'op>(
        &'log self,
        mut txn: SqlTransaction,
        log: &'op mut TransactionLogUpdates<'log>,
    ) -> Result<SqlTransaction, BookmarkTransactionError> {
        for (bookmark, log_entry) in self.force_deletes.iter() {
            log.push_log_entry(bookmark, log_entry);
            let (txn_, _) = DeleteBookmark::query_with_transaction(
                txn,
                &self.repo_id,
                bookmark.name(),
                bookmark.category(),
            )
            .await?;
            txn = txn_;
        }
        Ok(txn)
    }

    async fn store_deletes<'op, 'log: 'op>(
        &'log self,
        mut txn: SqlTransaction,
        log: &'op mut TransactionLogUpdates<'log>,
    ) -> Result<SqlTransaction, BookmarkTransactionError> {
        for (bookmark, old_cs_id, maybe_log_entry) in self.deletes.iter() {
            maybe_log_entry
                .as_ref()
                .map(|log_entry| log.push_log_entry(bookmark, log_entry));
            let (txn_, result) = DeleteBookmarkIf::query_with_transaction(
                txn,
                &self.repo_id,
                bookmark.name(),
                bookmark.category(),
                old_cs_id,
            )
            .await?;
            txn = txn_;
            if result.affected_rows() != 1 {
                return Err(BookmarkTransactionError::LogicError);
            }
        }
        Ok(txn)
    }

    async fn attempt_write(
        &self,
        txn: SqlTransaction,
    ) -> Result<SqlTransaction, BookmarkTransactionError> {
        let (mut txn, next_id) = Self::find_next_update_log_id(txn, self.repo_id).await?;

        let mut log = TransactionLogUpdates::new(next_id);

        txn = self.store_force_sets(txn, &mut log).await?;
        txn = self.store_creates(txn, &mut log).await?;
        txn = self.store_updates(txn, &mut log).await?;
        txn = self.store_force_deletes(txn, &mut log).await?;
        txn = self.store_deletes(txn, &mut log).await?;
        txn = self
            .store_log(txn, &log)
            .await
            .map_err(BookmarkTransactionError::RetryableError)?;

        Ok(txn)
    }
}

pub struct SqlBookmarksTransaction {
    write_connection: Connection,
    ctx: CoreContext,

    /// Bookmarks that have been seen already in this transaction.
    seen: HashSet<BookmarkKey>,

    /// Transaction updates.  A separate struct so that they can be
    /// moved into the future that will perform the database
    /// updates.
    payload: SqlBookmarksTransactionPayload,
}

impl SqlBookmarksTransaction {
    pub(crate) fn new(
        ctx: CoreContext,
        write_connection: Connection,
        repo_id: RepositoryId,
    ) -> Self {
        Self {
            write_connection,
            ctx,
            seen: HashSet::new(),
            payload: SqlBookmarksTransactionPayload::new(repo_id),
        }
    }

    pub fn check_not_seen(&mut self, bookmark: &BookmarkKey) -> Result<()> {
        if !self.seen.insert(bookmark.clone()) {
            return Err(anyhow!("{} bookmark was already used", bookmark));
        }
        Ok(())
    }
}

impl BookmarkTransaction for SqlBookmarksTransaction {
    fn update(
        &mut self,
        bookmark: &BookmarkKey,
        new_cs: ChangesetId,
        old_cs: ChangesetId,
        reason: BookmarkUpdateReason,
    ) -> Result<()> {
        self.check_not_seen(bookmark)?;
        let log = NewUpdateLogEntry::new(Some(old_cs), Some(new_cs), reason)?;
        self.payload.updates.push((
            bookmark.clone(),
            old_cs,
            new_cs,
            BookmarkKind::ALL_PUBLISHING,
            Some(log),
        ));
        Ok(())
    }

    fn update_scratch(
        &mut self,
        bookmark: &BookmarkKey,
        new_cs: ChangesetId,
        old_cs: ChangesetId,
    ) -> Result<()> {
        self.check_not_seen(bookmark)?;
        self.payload.updates.push((
            bookmark.clone(),
            old_cs,
            new_cs,
            &[BookmarkKind::Scratch],
            None, // Scratch bookmark updates are not logged.
        ));
        Ok(())
    }

    fn create(
        &mut self,
        bookmark: &BookmarkKey,
        new_cs: ChangesetId,
        reason: BookmarkUpdateReason,
    ) -> Result<()> {
        self.check_not_seen(bookmark)?;
        let log = NewUpdateLogEntry::new(None, Some(new_cs), reason)?;
        self.payload.creates.push((
            bookmark.clone(),
            new_cs,
            BookmarkKind::PullDefaultPublishing,
            Some(log),
        ));
        Ok(())
    }

    fn create_publishing(
        &mut self,
        bookmark: &BookmarkKey,
        new_cs: ChangesetId,
        reason: BookmarkUpdateReason,
    ) -> Result<()> {
        self.check_not_seen(bookmark)?;
        let log = NewUpdateLogEntry::new(None, Some(new_cs), reason)?;
        self.payload.creates.push((
            bookmark.clone(),
            new_cs,
            BookmarkKind::Publishing,
            Some(log),
        ));
        Ok(())
    }

    fn create_scratch(&mut self, bookmark: &BookmarkKey, new_cs: ChangesetId) -> Result<()> {
        self.check_not_seen(bookmark)?;
        self.payload.creates.push((
            bookmark.clone(),
            new_cs,
            BookmarkKind::Scratch,
            None, // Scratch bookmark updates are not logged.
        ));
        Ok(())
    }

    fn force_set(
        &mut self,
        bookmark: &BookmarkKey,
        new_cs: ChangesetId,
        reason: BookmarkUpdateReason,
    ) -> Result<()> {
        self.check_not_seen(bookmark)?;
        let log = NewUpdateLogEntry::new(None, Some(new_cs), reason)?;
        self.payload
            .force_sets
            .push((bookmark.clone(), new_cs, log));
        Ok(())
    }

    fn delete(
        &mut self,
        bookmark: &BookmarkKey,
        old_cs: ChangesetId,
        reason: BookmarkUpdateReason,
    ) -> Result<()> {
        self.check_not_seen(bookmark)?;
        let log = NewUpdateLogEntry::new(Some(old_cs), None, reason)?;
        self.payload
            .deletes
            .push((bookmark.clone(), old_cs, Some(log)));
        Ok(())
    }

    fn force_delete(&mut self, bookmark: &BookmarkKey, reason: BookmarkUpdateReason) -> Result<()> {
        self.check_not_seen(bookmark)?;
        let log = NewUpdateLogEntry::new(None, None, reason)?;
        self.payload.force_deletes.push((bookmark.clone(), log));
        Ok(())
    }

    fn delete_scratch(&mut self, bookmark: &BookmarkKey, old_cs: ChangesetId) -> Result<()> {
        self.check_not_seen(bookmark)?;
        self.payload.deletes.push((
            bookmark.clone(),
            old_cs,
            None, // Scratch bookmark updates are not logged.
        ));
        Ok(())
    }

    fn commit(self: Box<Self>) -> BoxFuture<'static, Result<bool>> {
        self.commit_with_hook(Arc::new(|_ctx, txn| future::ok(txn).boxed()))
    }

    /// commit_with_hook() can be used to have the same transaction to update two different database
    /// tables. `txn_hook()` should apply changes to the transaction.
    fn commit_with_hook(
        self: Box<Self>,
        txn_hook: BookmarkTransactionHook,
    ) -> BoxFuture<'static, Result<bool>> {
        let Self {
            ctx,
            payload,
            write_connection,
            ..
        } = *self;

        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlWrites);

        async move {
            let mut attempt = 0;
            let result = loop {
                attempt += 1;

                let mut txn = write_connection.start_transaction().await?;

                txn = match txn_hook(ctx.clone(), txn).await {
                    Ok(txn) => txn,
                    Err(BookmarkTransactionError::RetryableError(_))
                        if attempt < MAX_BOOKMARK_TRANSACTION_ATTEMPT_COUNT =>
                    {
                        continue;
                    }
                    err => break err,
                };

                match payload.attempt_write(txn).await {
                    Err(BookmarkTransactionError::RetryableError(_))
                        if attempt < MAX_BOOKMARK_TRANSACTION_ATTEMPT_COUNT =>
                    {
                        continue;
                    }
                    result => break result,
                }
            };

            // The number of `RetryableError`'s that were encountered
            let mut retryable_errors = attempt as i64 - 1;
            let result = match result {
                Ok(txn) => {
                    STATS::bookmarks_update_log_insert_success.add_value(1);
                    STATS::bookmarks_update_log_insert_success_attempt_count
                        .add_value(attempt as i64);
                    txn.commit().await?;
                    Ok(true)
                }
                Err(BookmarkTransactionError::LogicError) => {
                    // Logic error signifies that the transaction was rolled
                    // back, which likely means that bookmark has moved since
                    // our pushrebase finished. We need to retry the pushrebase
                    // Attempt count means one more than the number of `RetryableError`
                    // we hit before seeing this.
                    STATS::bookmarks_insert_logic_error.add_value(1);
                    STATS::bookmarks_insert_logic_error_attempt_count.add_value(attempt as i64);
                    Ok(false)
                }
                Err(BookmarkTransactionError::RetryableError(err)) => {
                    // Attempt count for `RetryableError` should always be equal
                    // to the MAX_BOOKMARK_TRANSACTION_ATTEMPT_COUNT, and hitting
                    // this error here basically means that this number of attempts
                    // was not enough, or the error was misclassified
                    STATS::bookmarks_insert_retryable_error.add_value(1);
                    STATS::bookmarks_insert_retryable_error_attempt_count.add_value(attempt as i64);
                    retryable_errors += 1;
                    Err(err)
                }
                Err(BookmarkTransactionError::Other(err)) => {
                    // `Other` error captures what we consider an "infrastructure"
                    // error, e.g. xdb went down during this transaction.
                    // Attempt count > 1 means the before we hit this error,
                    // we hit `RetryableError` a attempt count - 1 times.
                    STATS::bookmarks_insert_other_error.add_value(1);
                    STATS::bookmarks_insert_other_error_attempt_count.add_value(attempt as i64);
                    Err(err)
                }
            };
            STATS::bookmarks_update_log_insert_retry.add_value(retryable_errors);
            result
        }
        .boxed()
    }
}

#[cfg(test)]
pub(crate) async fn insert_bookmarks(
    conn: &Connection,
    rows: impl IntoIterator<Item = (&RepositoryId, &BookmarkKey, &ChangesetId, &BookmarkKind)>,
) -> Result<()> {
    let none = None;
    let rows = rows
        .into_iter()
        .map(|(r, b, c, k)| (r, &none, b.name(), b.category(), c, k))
        .collect::<Vec<_>>();
    InsertBookmarks::query(conn, rows.as_slice()).await?;
    Ok(())
}