/*
 * 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.
 */

//! Tests for the Changesets store.

use std::str::FromStr;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::sync::Arc;

use anyhow::Error;
use assert_matches::assert_matches;
use async_trait::async_trait;
use bonsai_hg_mapping::BonsaiHgMapping;
use bonsai_hg_mapping::BonsaiHgMappingEntry;
use bonsai_hg_mapping::BonsaiOrHgChangesetIds;
use bonsai_hg_mapping::CachingBonsaiHgMapping;
use bonsai_hg_mapping::ErrorKind;
use bonsai_hg_mapping::SqlBonsaiHgMappingBuilder;
use context::CoreContext;
use fbinit::FacebookInit;
use mercurial_types::HgChangesetId;
use mercurial_types::HgChangesetIdPrefix;
use mercurial_types::HgChangesetIdsResolvedFromPrefix;
use mercurial_types_mocks::nodehash as hg;
use mononoke_types::RepositoryId;
use mononoke_types_mocks::changesetid as bonsai;
use mononoke_types_mocks::repo::REPO_ZERO;
use rendezvous::RendezVousOptions;
use sql_construct::SqlConstruct;

async fn add_and_get<M: BonsaiHgMapping>(fb: FacebookInit, mapping: M) {
    let ctx = CoreContext::test_mock(fb);
    let entry = BonsaiHgMappingEntry {
        hg_cs_id: hg::ONES_CSID,
        bcs_id: bonsai::ONES_CSID,
    };
    assert!(
        mapping
            .add(&ctx, entry.clone())
            .await
            .expect("Adding new entry failed")
    );
    assert!(
        !mapping
            .add(&ctx, entry.clone())
            .await
            .expect("Adding same entry failed")
    );

    let result = mapping
        .get(&ctx, hg::ONES_CSID.into())
        .await
        .expect("Get failed");
    assert_eq!(result, vec![entry.clone()]);
    let result = mapping
        .get_hg_from_bonsai(&ctx, bonsai::ONES_CSID)
        .await
        .expect("Failed to get hg changeset by its bonsai counterpart");
    assert_eq!(result, Some(hg::ONES_CSID));
    let result = mapping
        .get_bonsai_from_hg(&ctx, hg::ONES_CSID)
        .await
        .expect("Failed to get bonsai changeset by its hg counterpart");
    assert_eq!(result, Some(bonsai::ONES_CSID));

    let same_bc_entry = BonsaiHgMappingEntry {
        hg_cs_id: hg::TWOS_CSID, // differ from entry.hg_cs_id
        bcs_id: bonsai::ONES_CSID,
    };
    let result = mapping
        .add(&ctx, same_bc_entry.clone())
        .await
        .expect_err("Conflicting entries should haved produced an error");
    assert_matches!(
        result.downcast::<ErrorKind>(),
        Ok(ErrorKind::ConflictingEntries(ref e0, ref e1)) if e0 == &entry && e1 == &same_bc_entry
    );

    let same_hg_entry = BonsaiHgMappingEntry {
        hg_cs_id: hg::ONES_CSID,
        bcs_id: bonsai::TWOS_CSID, // differ from entry.bcs_id
    };
    let result = mapping
        .add(&ctx, same_hg_entry.clone())
        .await
        .expect_err("Conflicting entries should haved produced an error");
    assert_matches!(
        result.downcast::<ErrorKind>(),
        Ok(ErrorKind::ConflictingEntries(ref e0, ref e1)) if e0 == &entry && e1 == &same_hg_entry
    );
}

async fn missing<M: BonsaiHgMapping>(fb: FacebookInit, mapping: M) {
    let ctx = CoreContext::test_mock(fb);
    let result = mapping
        .get(&ctx, bonsai::ONES_CSID.into())
        .await
        .expect("Failed to fetch missing changeset (should succeed with None instead)");
    assert_eq!(result, vec![]);
}

async fn get_many_hg_by_prefix<M: BonsaiHgMapping>(fb: FacebookInit, mapping: M) {
    let ctx = CoreContext::test_mock(fb);

    let entry1 = BonsaiHgMappingEntry {
        hg_cs_id: hg::ONES_CSID,
        bcs_id: bonsai::ONES_CSID,
    };
    let entry2 = BonsaiHgMappingEntry {
        hg_cs_id: hg::TWOS_CSID,
        bcs_id: bonsai::TWOS_CSID,
    };
    let entry3 = BonsaiHgMappingEntry {
        hg_cs_id: hg::FS_ES_CSID,
        bcs_id: bonsai::FS_ES_CSID,
    };
    let entry4 = BonsaiHgMappingEntry {
        hg_cs_id: hg::FS_CSID,
        bcs_id: bonsai::FS_CSID,
    };

    assert!(
        mapping
            .add(&ctx, entry1.clone())
            .await
            .expect("Adding entry1 failed")
    );
    assert!(
        mapping
            .add(&ctx, entry2.clone())
            .await
            .expect("Adding entry2 failed")
    );
    assert!(
        mapping
            .add(&ctx, entry3.clone())
            .await
            .expect("Adding entry3 failed")
    );
    assert!(
        mapping
            .add(&ctx, entry4.clone())
            .await
            .expect("Adding entry4 failed")
    );

    // found a single changeset
    let result = mapping
        .get_many_hg_by_prefix(
            &ctx,
            HgChangesetIdPrefix::from_bytes(&hg::ONES_CSID.as_ref()[0..8]).unwrap(),
            10,
        )
        .await
        .expect("Failed to get hg changeset by its prefix");

    assert_eq!(
        result,
        HgChangesetIdsResolvedFromPrefix::Single(hg::ONES_CSID)
    );

    // found a single changeset
    let result = mapping
        .get_many_hg_by_prefix(
            &ctx,
            HgChangesetIdPrefix::from_bytes(&hg::TWOS_CSID.as_ref()[0..10]).unwrap(),
            1,
        )
        .await
        .expect("Failed to get hg changeset by its prefix");

    assert_eq!(
        result,
        HgChangesetIdsResolvedFromPrefix::Single(hg::TWOS_CSID)
    );

    // found several changesets within the limit
    let result = mapping
        .get_many_hg_by_prefix(
            &ctx,
            HgChangesetIdPrefix::from_bytes(&hg::FS_CSID.as_ref()[0..8]).unwrap(),
            10,
        )
        .await
        .expect("Failed to get hg changeset by its prefix");

    assert_eq!(
        result,
        HgChangesetIdsResolvedFromPrefix::Multiple(vec![hg::FS_ES_CSID, hg::FS_CSID])
    );

    // found several changesets within the limit (try odd hex string prefix this time)
    let result = mapping
        .get_many_hg_by_prefix(&ctx, HgChangesetIdPrefix::from_str("fff").unwrap(), 10)
        .await
        .expect("Failed to get hg changeset by its prefix");

    assert_eq!(
        result,
        HgChangesetIdsResolvedFromPrefix::Multiple(vec![hg::FS_ES_CSID, hg::FS_CSID])
    );

    // found several changesets exceeding the limit
    let result = mapping
        .get_many_hg_by_prefix(
            &ctx,
            HgChangesetIdPrefix::from_bytes(&hg::FS_CSID.as_ref()[0..8]).unwrap(),
            1,
        )
        .await
        .expect("Failed to get hg changeset by its prefix");

    assert_eq!(
        result,
        HgChangesetIdsResolvedFromPrefix::TooMany(vec![hg::FS_ES_CSID])
    );

    // not found
    let result = mapping
        .get_many_hg_by_prefix(
            &ctx,
            HgChangesetIdPrefix::from_bytes(&hg::THREES_CSID.as_ref()[0..16]).unwrap(),
            10,
        )
        .await
        .expect("Failed to get hg changeset by its prefix");

    assert_eq!(result, HgChangesetIdsResolvedFromPrefix::NoMatch);
}

async fn get_hg_in_range<M: BonsaiHgMapping>(fb: FacebookInit, mapping: M) {
    let ctx = CoreContext::test_mock(fb);

    let entry1 = BonsaiHgMappingEntry {
        hg_cs_id: hg::ONES_CSID,
        bcs_id: bonsai::ONES_CSID,
    };
    let entry2 = BonsaiHgMappingEntry {
        hg_cs_id: hg::TWOS_CSID,
        bcs_id: bonsai::TWOS_CSID,
    };
    let entry3 = BonsaiHgMappingEntry {
        hg_cs_id: hg::THREES_CSID,
        bcs_id: bonsai::THREES_CSID,
    };

    for entry in [entry1, entry2, entry3] {
        assert!(mapping.add(&ctx, entry).await.unwrap());
    }

    assert!(
        mapping
            .get_hg_in_range(&ctx, hg::AS_CSID, hg::BS_CSID, 10)
            .await
            .unwrap()
            .is_empty(),
    );

    assert_eq!(
        vec![hg::ONES_CSID],
        mapping
            .get_hg_in_range(&ctx, hg::ONES_CSID, hg::ONES_CSID, 10)
            .await
            .unwrap()
    );

    assert_eq!(
        vec![hg::ONES_CSID, hg::TWOS_CSID],
        mapping
            .get_hg_in_range(&ctx, hg::ONES_CSID, hg::TWOS_CSID, 10)
            .await
            .unwrap()
    );

    assert_eq!(
        vec![hg::ONES_CSID],
        mapping
            .get_hg_in_range(&ctx, hg::ONES_CSID, hg::TWOS_CSID, 1)
            .await
            .unwrap()
    );

    assert_eq!(
        vec![hg::ONES_CSID, hg::TWOS_CSID, hg::THREES_CSID],
        mapping
            .get_hg_in_range(&ctx, hg::NULL_CSID, hg::FS_CSID, 10)
            .await
            .unwrap()
    );
}

struct CountedBonsaiHgMapping {
    mapping: Arc<dyn BonsaiHgMapping>,
    gets: Arc<AtomicUsize>,
    adds: Arc<AtomicUsize>,
    gets_many_hg_by_prefix: Arc<AtomicUsize>,
}

impl CountedBonsaiHgMapping {
    fn new(
        mapping: Arc<dyn BonsaiHgMapping>,
        gets: Arc<AtomicUsize>,
        adds: Arc<AtomicUsize>,
        gets_many_hg_by_prefix: Arc<AtomicUsize>,
    ) -> Self {
        Self {
            mapping,
            gets,
            adds,
            gets_many_hg_by_prefix,
        }
    }
}

#[async_trait]
impl BonsaiHgMapping for CountedBonsaiHgMapping {
    fn repo_id(&self) -> RepositoryId {
        self.mapping.repo_id()
    }

    async fn add(&self, ctx: &CoreContext, entry: BonsaiHgMappingEntry) -> Result<bool, Error> {
        self.adds.fetch_add(1, Ordering::Relaxed);
        self.mapping.add(ctx, entry).await
    }

    async fn get(
        &self,
        ctx: &CoreContext,
        cs_id: BonsaiOrHgChangesetIds,
    ) -> Result<Vec<BonsaiHgMappingEntry>, Error> {
        self.gets.fetch_add(1, Ordering::Relaxed);
        self.mapping.get(ctx, cs_id).await
    }

    async fn get_hg_in_range(
        &self,
        ctx: &CoreContext,
        low: HgChangesetId,
        high: HgChangesetId,
        limit: usize,
    ) -> Result<Vec<HgChangesetId>, Error> {
        self.gets_many_hg_by_prefix.fetch_add(1, Ordering::Relaxed);
        self.mapping.get_hg_in_range(ctx, low, high, limit).await
    }
}

async fn caching<M: BonsaiHgMapping + 'static>(fb: FacebookInit, mapping: M) {
    let ctx = CoreContext::test_mock(fb);
    let gets = Arc::new(AtomicUsize::new(0));
    let adds = Arc::new(AtomicUsize::new(0));
    let gets_many_hg_by_prefix = Arc::new(AtomicUsize::new(0));
    let mapping = CountedBonsaiHgMapping::new(
        Arc::new(mapping),
        gets.clone(),
        adds.clone(),
        gets_many_hg_by_prefix.clone(),
    );
    let mapping = CachingBonsaiHgMapping::new_test(Arc::new(mapping));

    let entry = BonsaiHgMappingEntry {
        hg_cs_id: hg::ONES_CSID,
        bcs_id: bonsai::ONES_CSID,
    };
    assert!(
        mapping
            .add(&ctx, entry.clone())
            .await
            .expect("Adding new entry failed")
    );

    let result = mapping
        .get_bonsai_from_hg(&ctx, hg::ONES_CSID)
        .await
        .expect("Failed to get bonsai changeset by its hg counterpart");
    assert_eq!(result, Some(bonsai::ONES_CSID));
    assert_eq!(gets.load(Ordering::Relaxed), 1);

    let result = mapping
        .get_bonsai_from_hg(&ctx, hg::ONES_CSID)
        .await
        .expect("Failed to get bonsai changeset by its hg counterpart");
    assert_eq!(result, Some(bonsai::ONES_CSID));
    assert_eq!(gets.load(Ordering::Relaxed), 1);

    let result = mapping
        .get_bonsai_from_hg(&ctx, hg::TWOS_CSID)
        .await
        .expect("Failed to get bonsai changeset by its hg counterpart");
    assert_eq!(result, None);
    assert_eq!(gets.load(Ordering::Relaxed), 2);
}

#[fbinit::test]
async fn test_add_and_get(fb: FacebookInit) {
    add_and_get(
        fb,
        SqlBonsaiHgMappingBuilder::with_sqlite_in_memory()
            .unwrap()
            .build(REPO_ZERO, RendezVousOptions::for_test()),
    )
    .await;
}

#[fbinit::test]
async fn test_missing(fb: FacebookInit) {
    missing(
        fb,
        SqlBonsaiHgMappingBuilder::with_sqlite_in_memory()
            .unwrap()
            .build(REPO_ZERO, RendezVousOptions::for_test()),
    )
    .await;
}

#[fbinit::test]
async fn test_caching(fb: FacebookInit) {
    caching(
        fb,
        SqlBonsaiHgMappingBuilder::with_sqlite_in_memory()
            .unwrap()
            .build(REPO_ZERO, RendezVousOptions::for_test()),
    )
    .await;
}

#[fbinit::test]
async fn test_get_many_hg_by_prefix(fb: FacebookInit) {
    get_many_hg_by_prefix(
        fb,
        SqlBonsaiHgMappingBuilder::with_sqlite_in_memory()
            .unwrap()
            .build(REPO_ZERO, RendezVousOptions::for_test()),
    )
    .await;
}

#[fbinit::test]
async fn test_get_hg_in_range(fb: FacebookInit) {
    get_hg_in_range(
        fb,
        SqlBonsaiHgMappingBuilder::with_sqlite_in_memory()
            .unwrap()
            .build(REPO_ZERO, RendezVousOptions::for_test()),
    )
    .await;
}

#[fbinit::test]
async fn test_overwrite(fb: FacebookInit) -> Result<(), Error> {
    let mapping = SqlBonsaiHgMappingBuilder::with_sqlite_in_memory()
        .unwrap()
        .with_overwrite()
        .build(REPO_ZERO, RendezVousOptions::for_test());

    let ctx = CoreContext::test_mock(fb);
    let entry = BonsaiHgMappingEntry {
        hg_cs_id: hg::ONES_CSID,
        bcs_id: bonsai::ONES_CSID,
    };

    assert!(
        mapping
            .add(&ctx, entry.clone())
            .await
            .expect("Adding new entry failed")
    );

    let entry = BonsaiHgMappingEntry {
        hg_cs_id: hg::ONES_CSID,
        bcs_id: bonsai::TWOS_CSID,
    };
    assert!(
        mapping
            .add(&ctx, entry.clone())
            .await
            .expect("Adding new entry failed")
    );

    let result = mapping.get(&ctx, hg::ONES_CSID.into()).await?;
    assert_eq!(result, vec![entry.clone()]);

    Ok(())
}


// @generated by autocargo
use std::env;
use std::fs;
use std::path::Path;

use thrift_compiler::Config;
use thrift_compiler::GenContext;

#[rustfmt::skip]
fn main() {
    // Rerun if this gets rewritten.
    println!("cargo:rerun-if-changed=thrift_build.rs");

    let out_dir = env::var_os("OUT_DIR").expect("OUT_DIR env not provided");
    let out_dir: &Path = out_dir.as_ref();
    fs::write(
        out_dir.join("cratemap"),
        "bonsai_svnrev_mapping crate
mononoke_types_thrift mononoke_types_thrift",
    ).expect("Failed to write cratemap");

    let conf = {
        let mut conf = Config::from_env(GenContext::Lib).expect("Failed to instantiate thrift_compiler::Config");

        let path_from_manifest_to_base: &Path = "../../../..".as_ref();
        let cargo_manifest_dir =
            env::var_os("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR not provided");
        let cargo_manifest_dir: &Path = cargo_manifest_dir.as_ref();
        let base_path = cargo_manifest_dir
            .join(path_from_manifest_to_base)
            .canonicalize()
            .expect("Failed to canonicalize base_path");
        // TODO: replace canonicalize() with std::path::absolute() when
        // https://github.com/rust-lang/rust/pull/91673 is available (~Rust 1.60)
        // and remove this block.
        #[cfg(windows)]
        let base_path = Path::new(
            base_path
                .as_path()
                .to_string_lossy()
                .trim_start_matches(r"\\?\"),
            )
            .to_path_buf();

        conf.base_path(base_path);

        let options = "types_crate=bonsai_svnrev_mapping_thrift__types";
        if !options.is_empty() {
            conf.options(options);
        }

        let include_srcs = vec![

        ];
        conf.include_srcs(include_srcs);

        conf
    };

    conf
        .run(&[
            "bonsai_svnrev_mapping.thrift"
        ])
        .expect("Failed while running thrift compilation");
}


// @generated by autocargo
::codegen_includer_proc_macro::include!();


/*
 * 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::collections::HashSet;
use std::sync::Arc;

use abomonation_derive::Abomonation;
use anyhow::anyhow;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use async_trait::async_trait;
use bonsai_svnrev_mapping_thrift as thrift;
use bytes::Bytes;
use caching_ext::get_or_fill;
use caching_ext::CacheDisposition;
use caching_ext::CacheHandlerFactory;
use caching_ext::CacheTtl;
use caching_ext::CachelibHandler;
use caching_ext::EntityStore;
use caching_ext::KeyedEntityStore;
use caching_ext::McErrorKind;
use caching_ext::McResult;
use caching_ext::MemcacheEntity;
use caching_ext::MemcacheHandler;
use context::CoreContext;
use fbthrift::compact_protocol;
use memcache::KeyGen;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use mononoke_types::Svnrev;

use super::BonsaiSvnrevMapping;
use super::BonsaiSvnrevMappingEntry;
use super::BonsaisOrSvnrevs;

#[derive(Abomonation, Clone, Debug, Eq, Hash, PartialEq)]
pub struct BonsaiSvnrevMappingCacheEntry {
    pub repo_id: RepositoryId,
    pub bcs_id: ChangesetId,
    pub svnrev: Svnrev,
}

impl BonsaiSvnrevMappingCacheEntry {
    pub fn new(repo_id: RepositoryId, bcs_id: ChangesetId, svnrev: Svnrev) -> Self {
        BonsaiSvnrevMappingCacheEntry {
            repo_id,
            bcs_id,
            svnrev,
        }
    }

    fn into_entry(self, repo_id: RepositoryId) -> Result<BonsaiSvnrevMappingEntry> {
        if self.repo_id == repo_id {
            Ok(BonsaiSvnrevMappingEntry {
                bcs_id: self.bcs_id,
                svnrev: self.svnrev,
            })
        } else {
            Err(anyhow!(
                "Cache returned invalid entry: repo {} returned for query to repo {}",
                self.repo_id,
                repo_id
            ))
        }
    }

    fn from_entry(entry: BonsaiSvnrevMappingEntry, repo_id: RepositoryId) -> Self {
        BonsaiSvnrevMappingCacheEntry {
            repo_id,
            bcs_id: entry.bcs_id,
            svnrev: entry.svnrev,
        }
    }
}

pub struct CachingBonsaiSvnrevMapping {
    cachelib: CachelibHandler<BonsaiSvnrevMappingCacheEntry>,
    memcache: MemcacheHandler,
    keygen: KeyGen,
    inner: Arc<dyn BonsaiSvnrevMapping>,
}

impl CachingBonsaiSvnrevMapping {
    pub fn new(
        inner: Arc<dyn BonsaiSvnrevMapping>,
        cache_handler_factory: CacheHandlerFactory,
    ) -> Self {
        Self {
            inner,
            cachelib: cache_handler_factory.cachelib(),
            memcache: cache_handler_factory.memcache(),
            keygen: Self::create_key_gen(),
        }
    }

    pub fn new_test(inner: Arc<dyn BonsaiSvnrevMapping>) -> Self {
        Self::new(inner, CacheHandlerFactory::Mocked)
    }

    fn create_key_gen() -> KeyGen {
        let key_prefix = "scm.mononoke.bonsai_svnrev_mapping";

        KeyGen::new(
            key_prefix,
            thrift::MC_CODEVER as u32,
            thrift::MC_SITEVER as u32,
        )
    }

    pub fn cachelib(&self) -> &CachelibHandler<BonsaiSvnrevMappingCacheEntry> {
        &self.cachelib
    }
}

#[async_trait]
impl BonsaiSvnrevMapping for CachingBonsaiSvnrevMapping {
    fn repo_id(&self) -> RepositoryId {
        self.inner.as_ref().repo_id()
    }

    async fn bulk_import(
        &self,
        ctx: &CoreContext,
        entries: &[BonsaiSvnrevMappingEntry],
    ) -> Result<(), Error> {
        self.inner.as_ref().bulk_import(ctx, entries).await
    }

    async fn get(
        &self,
        ctx: &CoreContext,
        objects: BonsaisOrSvnrevs,
    ) -> Result<Vec<BonsaiSvnrevMappingEntry>, Error> {
        let cache_request = (ctx, self);
        let repo_id = self.repo_id();

        let res = match objects {
            BonsaisOrSvnrevs::Bonsai(cs_ids) => {
                get_or_fill(&cache_request, cs_ids.into_iter().collect())
                    .await
                    .with_context(|| "Error fetching svnrevs via cache")?
                    .into_values()
                    .map(|val| val.into_entry(repo_id))
                    .collect::<Result<_>>()?
            }
            BonsaisOrSvnrevs::Svnrev(svnrevs) => {
                get_or_fill(&cache_request, svnrevs.into_iter().collect())
                    .await
                    .with_context(|| "Error fetching bonsais via cache")?
                    .into_values()
                    .map(|val| val.into_entry(repo_id))
                    .collect::<Result<_>>()?
            }
        };

        Ok(res)
    }
}

impl MemcacheEntity for BonsaiSvnrevMappingCacheEntry {
    fn serialize(&self) -> Bytes {
        let entry = thrift::BonsaiSvnrevMappingEntry {
            repo_id: self.repo_id.id(),
            bcs_id: self.bcs_id.into_thrift(),
            svnrev: self
                .svnrev
                .id()
                .try_into()
                .expect("Svnrevs must fit within a i64"),
        };
        compact_protocol::serialize(&entry)
    }

    fn deserialize(bytes: Bytes) -> McResult<Self> {
        let thrift::BonsaiSvnrevMappingEntry {
            repo_id,
            bcs_id,
            svnrev,
        } = compact_protocol::deserialize(bytes).map_err(|_| McErrorKind::Deserialization)?;

        let repo_id = RepositoryId::new(repo_id);
        let bcs_id = ChangesetId::from_thrift(bcs_id).map_err(|_| McErrorKind::Deserialization)?;
        let svnrev = Svnrev::new(
            svnrev
                .try_into()
                .map_err(|_| McErrorKind::Deserialization)?,
        );

        Ok(BonsaiSvnrevMappingCacheEntry {
            repo_id,
            bcs_id,
            svnrev,
        })
    }
}

type CacheRequest<'a> = (&'a CoreContext, &'a CachingBonsaiSvnrevMapping);

impl EntityStore<BonsaiSvnrevMappingCacheEntry> for CacheRequest<'_> {
    fn cachelib(&self) -> &CachelibHandler<BonsaiSvnrevMappingCacheEntry> {
        let (_, mapping) = self;
        &mapping.cachelib
    }

    fn keygen(&self) -> &KeyGen {
        let (_, mapping) = self;
        &mapping.keygen
    }

    fn memcache(&self) -> &MemcacheHandler {
        let (_, mapping) = self;
        &mapping.memcache
    }

    fn cache_determinator(&self, _: &BonsaiSvnrevMappingCacheEntry) -> CacheDisposition {
        CacheDisposition::Cache(CacheTtl::NoTtl)
    }

    caching_ext::impl_singleton_stats!("bonsai_svnrev_mapping");
}

#[async_trait]
impl KeyedEntityStore<ChangesetId, BonsaiSvnrevMappingCacheEntry> for CacheRequest<'_> {
    fn get_cache_key(&self, key: &ChangesetId) -> String {
        let (_, mapping) = self;
        format!("{}.bonsai.{}", mapping.repo_id(), key)
    }

    async fn get_from_db(
        &self,
        keys: HashSet<ChangesetId>,
    ) -> Result<HashMap<ChangesetId, BonsaiSvnrevMappingCacheEntry>, Error> {
        let (ctx, mapping) = self;
        let repo_id = mapping.repo_id();

        let res = mapping
            .inner
            .as_ref()
            .get(ctx, BonsaisOrSvnrevs::Bonsai(keys.into_iter().collect()))
            .await
            .with_context(|| "Error fetching svnrevs from bonsais from SQL")?;

        Result::<_, Error>::Ok(
            res.into_iter()
                .map(|e| {
                    (
                        e.bcs_id,
                        BonsaiSvnrevMappingCacheEntry::from_entry(e, repo_id),
                    )
                })
                .collect(),
        )
    }
}

#[async_trait]
impl KeyedEntityStore<Svnrev, BonsaiSvnrevMappingCacheEntry> for CacheRequest<'_> {
    fn get_cache_key(&self, key: &Svnrev) -> String {
        let (_, mapping) = self;
        format!("{}.svnrev.{}", mapping.repo_id(), key.id())
    }

    async fn get_from_db(
        &self,
        keys: HashSet<Svnrev>,
    ) -> Result<HashMap<Svnrev, BonsaiSvnrevMappingCacheEntry>, Error> {
        let (ctx, mapping) = self;
        let repo_id = mapping.repo_id();

        let res = mapping
            .inner
            .as_ref()
            .get(ctx, BonsaisOrSvnrevs::Svnrev(keys.into_iter().collect()))
            .await
            .with_context(|| "Error fetching bonsais from svnrevs from SQL")?;

        Result::<_, Error>::Ok(
            res.into_iter()
                .map(|e| {
                    (
                        e.svnrev,
                        BonsaiSvnrevMappingCacheEntry::from_entry(e, repo_id),
                    )
                })
                .collect(),
        )
    }
}


/*
 * 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.
 */

mod caching;
mod sql;

use anyhow::Error;
use async_trait::async_trait;
use context::CoreContext;
use mononoke_types::BonsaiChangeset;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use mononoke_types::Svnrev;
use slog::warn;

pub use crate::caching::CachingBonsaiSvnrevMapping;
pub use crate::sql::bulk_import_svnrevs;
pub use crate::sql::AddSvnrevsErrorKind;
pub use crate::sql::SqlBonsaiSvnrevMapping;
pub use crate::sql::SqlBonsaiSvnrevMappingBuilder;

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct BonsaiSvnrevMappingEntry {
    pub bcs_id: ChangesetId,
    pub svnrev: Svnrev,
}

impl BonsaiSvnrevMappingEntry {
    pub fn new(bcs_id: ChangesetId, svnrev: Svnrev) -> Self {
        BonsaiSvnrevMappingEntry { bcs_id, svnrev }
    }
}

pub enum BonsaisOrSvnrevs {
    Bonsai(Vec<ChangesetId>),
    Svnrev(Vec<Svnrev>),
}

impl BonsaisOrSvnrevs {
    pub fn is_empty(&self) -> bool {
        match self {
            BonsaisOrSvnrevs::Bonsai(v) => v.is_empty(),
            BonsaisOrSvnrevs::Svnrev(v) => v.is_empty(),
        }
    }
}

impl From<ChangesetId> for BonsaisOrSvnrevs {
    fn from(cs_id: ChangesetId) -> Self {
        BonsaisOrSvnrevs::Bonsai(vec![cs_id])
    }
}

impl From<Vec<ChangesetId>> for BonsaisOrSvnrevs {
    fn from(cs_ids: Vec<ChangesetId>) -> Self {
        BonsaisOrSvnrevs::Bonsai(cs_ids)
    }
}

impl From<Svnrev> for BonsaisOrSvnrevs {
    fn from(rev: Svnrev) -> Self {
        BonsaisOrSvnrevs::Svnrev(vec![rev])
    }
}

impl From<Vec<Svnrev>> for BonsaisOrSvnrevs {
    fn from(revs: Vec<Svnrev>) -> Self {
        BonsaisOrSvnrevs::Svnrev(revs)
    }
}

#[facet::facet]
#[async_trait]
pub trait BonsaiSvnrevMapping: Send + Sync {
    fn repo_id(&self) -> RepositoryId;

    async fn bulk_import(
        &self,
        ctx: &CoreContext,
        entries: &[BonsaiSvnrevMappingEntry],
    ) -> Result<(), Error>;

    async fn get(
        &self,
        ctx: &CoreContext,
        field: BonsaisOrSvnrevs,
    ) -> Result<Vec<BonsaiSvnrevMappingEntry>, Error>;

    async fn get_svnrev_from_bonsai(
        &self,
        ctx: &CoreContext,
        bcs_id: ChangesetId,
    ) -> Result<Option<Svnrev>, Error> {
        let result = self
            .get(ctx, BonsaisOrSvnrevs::Bonsai(vec![bcs_id]))
            .await?;
        Ok(result.into_iter().next().map(|entry| entry.svnrev))
    }

    async fn get_bonsai_from_svnrev(
        &self,
        ctx: &CoreContext,
        svnrev: Svnrev,
    ) -> Result<Option<ChangesetId>, Error> {
        let result = self
            .get(ctx, BonsaisOrSvnrevs::Svnrev(vec![svnrev]))
            .await?;
        Ok(result.into_iter().next().map(|entry| entry.bcs_id))
    }

    async fn bulk_import_from_bonsai(
        &self,
        ctx: &CoreContext,
        changesets: &[BonsaiChangeset],
    ) -> anyhow::Result<()> {
        let mut entries = vec![];
        for bcs in changesets.iter() {
            match Svnrev::from_bcs(bcs) {
                Ok(svnrev) => {
                    let entry = BonsaiSvnrevMappingEntry::new(bcs.get_changeset_id(), svnrev);
                    entries.push(entry);
                }
                Err(e) => {
                    warn!(ctx.logger(), "Couldn't fetch svnrev from commit: {:?}", e);
                }
            }
        }
        self.bulk_import(ctx, &entries).await?;
        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 ::sql::Connection;
use ::sql_ext::mononoke_queries;
use anyhow::Error;
use async_trait::async_trait;
use context::CoreContext;
use context::PerfCounterType;
use mononoke_types::BonsaiChangeset;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use mononoke_types::Svnrev;
use slog::warn;
use sql_construct::SqlConstruct;
use sql_construct::SqlConstructFromMetadataDatabaseConfig;
use sql_ext::SqlConnections;
use thiserror::Error;

use super::BonsaiSvnrevMapping;
use super::BonsaiSvnrevMappingEntry;
use super::BonsaisOrSvnrevs;

mononoke_queries! {
    write DangerouslyAddSvnrevs(values: (
        repo_id: RepositoryId,
        bcs_id: ChangesetId,
        svnrev: Svnrev,
    )) {
        insert_or_ignore,
        "{insert_or_ignore} INTO bonsai_svnrev_mapping (repo_id, bcs_id, svnrev) VALUES {values}"
    }

    read SelectMappingByBonsai(
        repo_id: RepositoryId,
        >list bcs_id: ChangesetId
    ) -> (ChangesetId, Svnrev) {
        "SELECT bcs_id, svnrev
         FROM bonsai_svnrev_mapping
         WHERE repo_id = {repo_id} AND bcs_id in {bcs_id}"
    }

    read SelectMappingBySvnrev(
        repo_id: RepositoryId,
        >list svnrev: Svnrev
    ) -> (ChangesetId, Svnrev) {
        "SELECT bcs_id, svnrev
         FROM bonsai_svnrev_mapping
         WHERE repo_id = {repo_id} AND svnrev in {svnrev}"
    }
}

pub struct SqlBonsaiSvnrevMapping {
    connections: SqlConnections,
    repo_id: RepositoryId,
}

#[derive(Clone)]
pub struct SqlBonsaiSvnrevMappingBuilder {
    connections: SqlConnections,
}

impl SqlConstruct for SqlBonsaiSvnrevMappingBuilder {
    const LABEL: &'static str = "bonsai_svnrev_mapping";

    const CREATION_QUERY: &'static str =
        include_str!("../schemas/sqlite-bonsai-svnrev-mapping.sql");

    fn from_sql_connections(connections: SqlConnections) -> Self {
        Self { connections }
    }
}

impl SqlConstructFromMetadataDatabaseConfig for SqlBonsaiSvnrevMappingBuilder {}

impl SqlBonsaiSvnrevMappingBuilder {
    pub fn build(self, repo_id: RepositoryId) -> SqlBonsaiSvnrevMapping {
        SqlBonsaiSvnrevMapping {
            connections: self.connections,
            repo_id,
        }
    }
}

#[async_trait]
impl BonsaiSvnrevMapping for SqlBonsaiSvnrevMapping {
    fn repo_id(&self) -> RepositoryId {
        self.repo_id
    }

    async fn bulk_import(
        &self,
        ctx: &CoreContext,
        entries: &[BonsaiSvnrevMappingEntry],
    ) -> Result<(), Error> {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlWrites);

        let entries: Vec<_> = entries
            .iter()
            .map(|entry| (&self.repo_id, &entry.bcs_id, &entry.svnrev))
            .collect();

        DangerouslyAddSvnrevs::query(&self.connections.write_connection, &entries[..]).await?;

        Ok(())
    }

    async fn get(
        &self,
        ctx: &CoreContext,
        objects: BonsaisOrSvnrevs,
    ) -> Result<Vec<BonsaiSvnrevMappingEntry>, Error> {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);

        let mut mappings =
            select_mapping(&self.connections.read_connection, self.repo_id, &objects).await?;

        let left_to_fetch = filter_fetched_objects(objects, &mappings[..]);

        if left_to_fetch.is_empty() {
            return Ok(mappings);
        }

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

        let mut master_mappings = select_mapping(
            &self.connections.read_master_connection,
            self.repo_id,
            &left_to_fetch,
        )
        .await?;
        mappings.append(&mut master_mappings);
        Ok(mappings)
    }
}

fn filter_fetched_objects(
    objects: BonsaisOrSvnrevs,
    mappings: &[BonsaiSvnrevMappingEntry],
) -> BonsaisOrSvnrevs {
    match objects {
        BonsaisOrSvnrevs::Bonsai(cs_ids) => {
            let bcs_fetched: HashSet<_> = mappings.iter().map(|m| &m.bcs_id).collect();

            BonsaisOrSvnrevs::Bonsai(
                cs_ids
                    .iter()
                    .filter_map(|cs| {
                        if !bcs_fetched.contains(cs) {
                            Some(*cs)
                        } else {
                            None
                        }
                    })
                    .collect(),
            )
        }
        BonsaisOrSvnrevs::Svnrev(svnrevs) => {
            let svnrevs_fetched: HashSet<_> = mappings.iter().map(|m| &m.svnrev).collect();

            BonsaisOrSvnrevs::Svnrev(
                svnrevs
                    .iter()
                    .filter_map(|svnrev| {
                        if !svnrevs_fetched.contains(svnrev) {
                            Some(*svnrev)
                        } else {
                            None
                        }
                    })
                    .collect(),
            )
        }
    }
}

async fn select_mapping(
    connection: &Connection,
    repo_id: RepositoryId,
    objects: &BonsaisOrSvnrevs,
) -> Result<Vec<BonsaiSvnrevMappingEntry>, Error> {
    if objects.is_empty() {
        return Ok(vec![]);
    }

    let rows = match objects {
        BonsaisOrSvnrevs::Bonsai(bcs_ids) => {
            SelectMappingByBonsai::query(connection, &repo_id, &bcs_ids[..]).await?
        }
        BonsaisOrSvnrevs::Svnrev(svnrevs) => {
            SelectMappingBySvnrev::query(connection, &repo_id, &svnrevs[..]).await?
        }
    };

    Ok(rows
        .into_iter()
        .map(move |(bcs_id, svnrev)| BonsaiSvnrevMappingEntry { bcs_id, svnrev })
        .collect())
}

/// This method is for importing Svnrevs in bulk from a set of BonsaiChangesets where you know
/// they are correct.
pub async fn bulk_import_svnrevs<'a>(
    ctx: &'a CoreContext,
    svnrevs_store: &'a impl BonsaiSvnrevMapping,
    changesets: impl IntoIterator<Item = &'a BonsaiChangeset>,
) -> Result<(), Error> {
    let mut entries = vec![];
    for bcs in changesets.into_iter() {
        match Svnrev::from_bcs(bcs) {
            Ok(svnrev) => {
                let entry = BonsaiSvnrevMappingEntry::new(bcs.get_changeset_id(), svnrev);
                entries.push(entry);
            }
            Err(e) => {
                warn!(ctx.logger(), "Couldn't fetch svnrev from commit: {:?}", e);
            }
        }
    }

    svnrevs_store.bulk_import(ctx, &entries).await?;

    Ok(())
}

#[derive(Debug, Error)]
pub enum AddSvnrevsErrorKind {
    #[error("Conflict detected while inserting Svnrevs")]
    Conflict,

    #[error("Internal error occurred while inserting Svnrevs")]
    InternalError(#[from] Error),
}


/*
 * 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::sync::Arc;

use anyhow::Error;
use bonsai_svnrev_mapping::BonsaiSvnrevMapping;
use bonsai_svnrev_mapping::BonsaiSvnrevMappingEntry;
use bonsai_svnrev_mapping::BonsaisOrSvnrevs;
use bonsai_svnrev_mapping::CachingBonsaiSvnrevMapping;
use bonsai_svnrev_mapping::SqlBonsaiSvnrevMappingBuilder;
use context::CoreContext;
use fbinit::FacebookInit;
use mononoke_types_mocks::changesetid as bonsai;
use mononoke_types_mocks::repo::REPO_ZERO;
use mononoke_types_mocks::svnrev::*;
use sql_construct::SqlConstruct;

#[fbinit::test]
async fn test_add_and_get(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let mapping = SqlBonsaiSvnrevMappingBuilder::with_sqlite_in_memory()?.build(REPO_ZERO);

    let entry = BonsaiSvnrevMappingEntry {
        bcs_id: bonsai::ONES_CSID,
        svnrev: SVNREV_ZERO,
    };

    mapping.bulk_import(&ctx, &[entry.clone()]).await?;

    let result = mapping
        .get(&ctx, BonsaisOrSvnrevs::Bonsai(vec![bonsai::ONES_CSID]))
        .await?;
    assert_eq!(result, vec![entry.clone()]);
    let result = mapping
        .get_svnrev_from_bonsai(&ctx, bonsai::ONES_CSID)
        .await?;
    assert_eq!(result, Some(SVNREV_ZERO));
    let result = mapping.get_bonsai_from_svnrev(&ctx, SVNREV_ZERO).await?;
    assert_eq!(result, Some(bonsai::ONES_CSID));

    Ok(())
}

#[fbinit::test]
async fn test_bulk_import(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let mapping = SqlBonsaiSvnrevMappingBuilder::with_sqlite_in_memory()?.build(REPO_ZERO);

    let entry1 = BonsaiSvnrevMappingEntry {
        bcs_id: bonsai::ONES_CSID,
        svnrev: SVNREV_ZERO,
    };
    let entry2 = BonsaiSvnrevMappingEntry {
        bcs_id: bonsai::TWOS_CSID,
        svnrev: SVNREV_ONE,
    };
    let entry3 = BonsaiSvnrevMappingEntry {
        bcs_id: bonsai::THREES_CSID,
        svnrev: SVNREV_TWO,
    };

    mapping
        .bulk_import(&ctx, &[entry1.clone(), entry2.clone(), entry3.clone()])
        .await?;

    Ok(())
}

#[fbinit::test]
async fn test_missing(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let mapping = SqlBonsaiSvnrevMappingBuilder::with_sqlite_in_memory()?.build(REPO_ZERO);

    let result = mapping
        .get(&ctx, BonsaisOrSvnrevs::Bonsai(vec![bonsai::ONES_CSID]))
        .await?;

    assert_eq!(result, vec![]);

    Ok(())
}

#[fbinit::test]
async fn test_caching(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let mapping =
        Arc::new(SqlBonsaiSvnrevMappingBuilder::with_sqlite_in_memory()?.build(REPO_ZERO));
    let caching = CachingBonsaiSvnrevMapping::new_test(mapping.clone());

    let store = caching
        .cachelib()
        .mock_store()
        .expect("new_test gives us a MockStore");

    let e0 = BonsaiSvnrevMappingEntry {
        bcs_id: bonsai::ONES_CSID,
        svnrev: SVNREV_ONE,
    };

    let e1 = BonsaiSvnrevMappingEntry {
        bcs_id: bonsai::TWOS_CSID,
        svnrev: SVNREV_TWO,
    };

    mapping.bulk_import(&ctx, &[e0, e1]).await?;

    assert_eq!(
        caching
            .get_svnrev_from_bonsai(&ctx, bonsai::ONES_CSID)
            .await?,
        Some(SVNREV_ONE)
    );

    assert_eq!(store.stats().gets, 1);
    assert_eq!(store.stats().hits, 0);
    assert_eq!(store.stats().sets, 1);

    assert_eq!(
        caching
            .get_svnrev_from_bonsai(&ctx, bonsai::ONES_CSID)
            .await?,
        Some(SVNREV_ONE)
    );

    assert_eq!(store.stats().gets, 2);
    assert_eq!(store.stats().hits, 1);
    assert_eq!(store.stats().sets, 1);

    assert_eq!(
        caching
            .get_svnrev_from_bonsai(&ctx, bonsai::TWOS_CSID)
            .await?,
        Some(SVNREV_TWO)
    );

    assert_eq!(
        caching.get_bonsai_from_svnrev(&ctx, SVNREV_ONE).await?,
        Some(bonsai::ONES_CSID)
    );

    assert_eq!(
        caching.get_bonsai_from_svnrev(&ctx, SVNREV_TWO).await?,
        Some(bonsai::TWOS_CSID)
    );

    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.
 */

mod sql;

use anyhow::Result;
use async_trait::async_trait;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;

pub use crate::sql::SqlBonsaiTagMapping;
pub use crate::sql::SqlBonsaiTagMappingBuilder;

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct BonsaiTagMappingEntry {
    pub changeset_id: ChangesetId,
    pub tag_name: String,
}

impl BonsaiTagMappingEntry {
    pub fn new(changeset_id: ChangesetId, tag_name: String) -> Self {
        BonsaiTagMappingEntry {
            changeset_id,
            tag_name,
        }
    }
}

#[facet::facet]
#[async_trait]
/// Facet trait representing Bonsai Changeset to Git Tag mapping
pub trait BonsaiTagMapping: Send + Sync {
    /// The repository for which this mapping has been created
    fn repo_id(&self) -> RepositoryId;

    /// Fetch the changeset id correponding to the tag name in the
    /// given repo, if one exists
    async fn get_changeset_by_tag_name(&self, tag_name: String) -> Result<Option<ChangesetId>>;

    /// Fetch the tag names corresponding to the input changeset id
    /// for the given repo. Note that there can potentially be multiple
    /// tags that map to the same changeset
    async fn get_tag_names_by_changeset(
        &self,
        changeset_id: ChangesetId,
    ) -> Result<Option<Vec<String>>>;

    /// Add new tag name to bonsai changeset mappings
    async fn add_mappings(&self, entries: Vec<BonsaiTagMappingEntry>) -> Result<()>;
}


/*
 * 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 ::sql_ext::mononoke_queries;
use anyhow::Context;
use anyhow::Result;
use async_trait::async_trait;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use sql_construct::SqlConstruct;
use sql_construct::SqlConstructFromMetadataDatabaseConfig;
use sql_ext::SqlConnections;

use super::BonsaiTagMapping;
use super::BonsaiTagMappingEntry;

mononoke_queries! {
    write AddBonsaiTagMapping(values: (
        repo_id: RepositoryId,
        tag_name: String,
        changeset_id: ChangesetId,
    )) {
        insert_or_ignore,
        "{insert_or_ignore} INTO bonsai_tag_mapping (repo_id, tag_name, changeset_id) VALUES {values}"
    }

    read SelectMappingByChangeset(
        repo_id: RepositoryId,
        changeset_id: ChangesetId
    ) -> (String, ChangesetId) {
        "SELECT tag_name, changeset_id
         FROM bonsai_tag_mapping
         WHERE repo_id = {repo_id} AND changeset_id = {changeset_id}"
    }

    read SelectMappingByTagName(
        repo_id: RepositoryId,
        tag_name: String,
    ) -> (String, ChangesetId) {
        "SELECT tag_name, changeset_id
         FROM bonsai_tag_mapping
         WHERE repo_id = {repo_id} AND tag_name = {tag_name}"
    }
}

pub struct SqlBonsaiTagMapping {
    connections: SqlConnections,
    repo_id: RepositoryId,
}

#[derive(Clone)]
pub struct SqlBonsaiTagMappingBuilder {
    connections: SqlConnections,
}

impl SqlConstruct for SqlBonsaiTagMappingBuilder {
    const LABEL: &'static str = "bonsai_tag_mapping";

    const CREATION_QUERY: &'static str = include_str!("../schemas/sqlite-bonsai-tag-mapping.sql");

    fn from_sql_connections(connections: SqlConnections) -> Self {
        Self { connections }
    }
}

impl SqlConstructFromMetadataDatabaseConfig for SqlBonsaiTagMappingBuilder {}

impl SqlBonsaiTagMappingBuilder {
    pub fn build(self, repo_id: RepositoryId) -> SqlBonsaiTagMapping {
        SqlBonsaiTagMapping {
            connections: self.connections,
            repo_id,
        }
    }
}

#[async_trait]
impl BonsaiTagMapping for SqlBonsaiTagMapping {
    fn repo_id(&self) -> RepositoryId {
        self.repo_id
    }

    async fn get_changeset_by_tag_name(&self, tag_name: String) -> Result<Option<ChangesetId>> {
        let results = SelectMappingByTagName::query(
            &self.connections.read_connection,
            &self.repo_id,
            &tag_name,
        )
        .await
        .with_context(|| {
            format!(
                "Failure in fetching changeset for tag {} in repo {}",
                tag_name, self.repo_id
            )
        })?;
        // This should not happen but since this is new code, extra checks dont hurt.
        if results.len() > 1 {
            anyhow::bail!(
                "Multiple changesets returned for tag {} in repo {}",
                tag_name,
                self.repo_id
            )
        }
        Ok(results
            .into_iter()
            .next()
            .map(|(_, changeset_id)| changeset_id))
    }

    async fn get_tag_names_by_changeset(
        &self,
        changeset_id: ChangesetId,
    ) -> Result<Option<Vec<String>>> {
        let results = SelectMappingByChangeset::query(
            &self.connections.read_connection,
            &self.repo_id,
            &changeset_id,
        )
        .await
        .with_context(|| {
            format!(
                "Failure in fetching tag for changeset {:?} in repo {}",
                changeset_id, self.repo_id
            )
        })?;

        let values = results
            .into_iter()
            .map(|(tag_name, _)| tag_name)
            .collect::<Vec<_>>();
        let output = (!values.is_empty()).then_some(values);
        return Ok(output);
    }

    async fn add_mappings(&self, entries: Vec<BonsaiTagMappingEntry>) -> Result<()> {
        let entries: Vec<_> = entries
            .iter()
            .map(|entry| (&self.repo_id, &entry.tag_name, &entry.changeset_id))
            .collect();
        AddBonsaiTagMapping::query(&self.connections.write_connection, entries.as_slice())
            .await
            .with_context(|| {
                format!(
                    "Failed to add mappings in repo {} for entries {:?}",
                    self.repo_id, entries,
                )
            })?;
        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 anyhow::Error;
use bonsai_tag_mapping::BonsaiTagMapping;
use bonsai_tag_mapping::BonsaiTagMappingEntry;
use bonsai_tag_mapping::SqlBonsaiTagMappingBuilder;
use fbinit::FacebookInit;
use mononoke_types_mocks::changesetid as bonsai;
use mononoke_types_mocks::repo::REPO_ZERO;
use sql_construct::SqlConstruct;

#[fbinit::test]
async fn test_add_and_get(_: FacebookInit) -> Result<(), Error> {
    let mapping = SqlBonsaiTagMappingBuilder::with_sqlite_in_memory()?.build(REPO_ZERO);
    let tag_name = "JustATag";
    let entry = BonsaiTagMappingEntry {
        changeset_id: bonsai::ONES_CSID,
        tag_name: tag_name.to_string(),
    };

    mapping.add_mappings(vec![entry.clone()]).await?;

    let result = mapping
        .get_changeset_by_tag_name(entry.tag_name.clone())
        .await?;
    assert_eq!(result, Some(bonsai::ONES_CSID));

    let result = mapping
        .get_tag_names_by_changeset(bonsai::ONES_CSID)
        .await?;
    assert_eq!(result, Some(vec![tag_name.to_string()]));
    Ok(())
}

#[fbinit::test]
async fn test_get_without_add(_: FacebookInit) -> Result<(), Error> {
    let mapping = SqlBonsaiTagMappingBuilder::with_sqlite_in_memory()?.build(REPO_ZERO);
    let result = mapping
        .get_changeset_by_tag_name("JustATag".to_string())
        .await?;
    assert_eq!(result, None);

    let result = mapping
        .get_tag_names_by_changeset(bonsai::ONES_CSID)
        .await?;
    assert_eq!(result, None);
    Ok(())
}

#[fbinit::test]
async fn test_get_multiple_tags(_: FacebookInit) -> Result<(), Error> {
    let mapping = SqlBonsaiTagMappingBuilder::with_sqlite_in_memory()?.build(REPO_ZERO);
    let entry = BonsaiTagMappingEntry {
        changeset_id: bonsai::ONES_CSID,
        tag_name: "JustATag".to_string(),
    };
    let another_entry = BonsaiTagMappingEntry {
        changeset_id: bonsai::ONES_CSID,
        tag_name: "AnotherTag".to_string(),
    };
    mapping.add_mappings(vec![entry, another_entry]).await?;

    let result = mapping
        .get_tag_names_by_changeset(bonsai::ONES_CSID)
        .await?
        .expect("None tags returned for the input changeset");
    assert_eq!(
        HashSet::from_iter(result),
        HashSet::from(["JustATag".to_string(), "AnotherTag".to_string()])
    );
    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::HashMap;
use std::collections::HashSet;

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use blobstore::Loadable;
use bookmarks::BookmarkUpdateReason;
use bookmarks_types::BookmarkKey;
use bookmarks_types::BookmarkKind;
use bytes::Bytes;
use context::CoreContext;
use cross_repo_sync::CHANGE_XREPO_MAPPING_EXTRA;
use futures::future;
use futures::stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use futures_ext::FbStreamExt;
use hooks::CrossRepoPushSource;
use hooks::HookManager;
use hooks::PushAuthoredBy;
use mononoke_types::BonsaiChangeset;
use mononoke_types::ChangesetId;
use repo_authorization::AuthorizationContext;
use skeleton_manifest::RootSkeletonManifestId;
use tunables::tunables;

use crate::hook_running::run_hooks;
use crate::restrictions::should_run_hooks;
use crate::BookmarkMovementError;
use crate::Repo;

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub(crate) enum AdditionalChangesets {
    None,
    Ancestors(ChangesetId),
    Range {
        head: ChangesetId,
        base: ChangesetId,
    },
}

pub(crate) struct AffectedChangesets {
    /// Changesets that are being added to the repository and to this bookmark.
    new_changesets: HashMap<ChangesetId, BonsaiChangeset>,

    /// Changesets that are being used as a source for pushrebase.
    source_changesets: HashSet<BonsaiChangeset>,

    /// Additional changesets, if they have been loaded.
    additional_changesets: Option<HashSet<BonsaiChangeset>>,
}

impl AffectedChangesets {
    pub(crate) fn new() -> Self {
        Self {
            new_changesets: HashMap::new(),
            source_changesets: HashSet::new(),
            additional_changesets: None,
        }
    }

    pub(crate) fn with_source_changesets(source_changesets: HashSet<BonsaiChangeset>) -> Self {
        Self {
            new_changesets: HashMap::new(),
            source_changesets,
            additional_changesets: None,
        }
    }

    pub(crate) fn add_new_changesets(
        &mut self,
        new_changesets: HashMap<ChangesetId, BonsaiChangeset>,
    ) {
        if self.new_changesets.is_empty() {
            self.new_changesets = new_changesets;
        } else {
            self.new_changesets.extend(new_changesets);
        }
    }

    pub(crate) fn new_changesets(&self) -> &HashMap<ChangesetId, BonsaiChangeset> {
        &self.new_changesets
    }

    pub(crate) fn source_changesets(&self) -> &HashSet<BonsaiChangeset> {
        &self.source_changesets
    }

    fn adding_new_changesets_to_repo(&self) -> bool {
        !self.source_changesets.is_empty() || !self.new_changesets.is_empty()
    }

    /// Load bonsais in the additional changeset range that are not already in
    /// `new_changesets` and are ancestors of `head` but not ancestors of `base`
    /// or any of the `hooks_skip_ancestors_of` bookmarks for the named
    /// bookmark.
    ///
    /// These are the additional bonsais that we need to run hooks on for
    /// bookmark moves.
    async fn load_additional_changesets(
        &mut self,
        ctx: &CoreContext,
        repo: &impl Repo,
        bookmark: &BookmarkKey,
        additional_changesets: AdditionalChangesets,
    ) -> Result<(), Error> {
        if self.additional_changesets.is_some() {
            return Ok(());
        }

        let (head, base) = match additional_changesets {
            AdditionalChangesets::None => {
                self.additional_changesets = Some(HashSet::new());
                return Ok(());
            }
            AdditionalChangesets::Ancestors(head) => (head, None),
            AdditionalChangesets::Range { head, base } => (head, Some(base)),
        };

        let mut exclude_bookmarks: HashSet<_> = repo
            .repo_bookmark_attrs()
            .select(bookmark)
            .flat_map(|attr| attr.params().hooks_skip_ancestors_of.iter())
            .cloned()
            .collect();
        exclude_bookmarks.remove(bookmark);

        let mut excludes: HashSet<_> = stream::iter(exclude_bookmarks)
            .map(|bookmark| repo.bookmarks().get(ctx.clone(), &bookmark))
            .buffered(100)
            .try_filter_map(|maybe_cs_id| async move { Ok(maybe_cs_id) })
            .try_collect()
            .await?;
        excludes.extend(base);

        let range = repo
            .commit_graph()
            .ancestors_difference_stream(ctx, vec![head], excludes.into_iter().collect())
            .await?
            .yield_periodically()
            .try_filter(|bcs_id| {
                let exists = self.new_changesets.contains_key(bcs_id);
                future::ready(!exists)
            });

        let limit = match tunables()
            .hooks_additional_changesets_limit()
            .unwrap_or_default()
        {
            limit if limit > 0 => limit as usize,
            _ => std::usize::MAX,
        };

        let additional_changesets = if tunables()
            .run_hooks_on_additional_changesets()
            .unwrap_or_default()
        {
            let bonsais = range
                .and_then({
                    let mut count = 0;
                    move |bcs_id| {
                        count += 1;
                        if count > limit {
                            future::ready(Err(anyhow!(
                                "hooks additional changesets limit reached at {}",
                                bcs_id
                            )))
                        } else {
                            future::ready(Ok(bcs_id))
                        }
                    }
                })
                .map(|res| async move {
                    match res {
                        Ok(bcs_id) => Ok(bcs_id.load(ctx, repo.repo_blobstore()).await?),
                        Err(e) => Err(e),
                    }
                })
                .buffered(100)
                .try_collect::<HashSet<_>>()
                .await?;

            ctx.scuba()
                .clone()
                .add("hook_running_additional_changesets", bonsais.len())
                .log_with_msg("Running hooks for additional changesets", None);

            bonsais
        } else {
            // Logging-only mode.  Work out how many changesets we would have run
            // on, and whether the limit would have been reached.
            let count = range
                .take(limit)
                .try_fold(0usize, |acc, _| async move { Ok(acc + 1) })
                .await?;

            let mut scuba = ctx.scuba().clone();
            scuba.add("hook_running_additional_changesets", count);
            if count >= limit {
                scuba.add("hook_running_additional_changesets_limit_reached", true);
            }
            scuba.log_with_msg("Hook running skipping additional changesets", None);
            HashSet::new()
        };

        self.additional_changesets = Some(additional_changesets);
        Ok(())
    }

    fn is_empty(&self) -> bool {
        self.new_changesets.is_empty()
            && self.source_changesets.is_empty()
            && self
                .additional_changesets
                .as_ref()
                .map_or(true, HashSet::is_empty)
    }

    fn iter(&self) -> impl Iterator<Item = &BonsaiChangeset> + Clone {
        self.new_changesets
            .values()
            .chain(self.source_changesets.iter())
            .chain(self.additional_changesets.iter().flatten())
    }

    /// Check all applicable restrictions on the affected changesets.
    pub(crate) async fn check_restrictions(
        &mut self,
        ctx: &CoreContext,
        authz: &AuthorizationContext,
        repo: &impl Repo,
        hook_manager: &HookManager,
        bookmark: &BookmarkKey,
        pushvars: Option<&HashMap<String, Bytes>>,
        reason: BookmarkUpdateReason,
        kind: BookmarkKind,
        additional_changesets: AdditionalChangesets,
        cross_repo_push_source: CrossRepoPushSource,
    ) -> Result<(), BookmarkMovementError> {
        self.check_extras(ctx, repo, bookmark, kind, additional_changesets)
            .await?;

        self.check_case_conflicts(ctx, repo, bookmark, kind, additional_changesets)
            .await?;

        self.check_hooks(
            ctx,
            authz,
            repo,
            hook_manager,
            bookmark,
            pushvars,
            reason,
            kind,
            additional_changesets,
            cross_repo_push_source,
        )
        .await?;

        self.check_path_permissions(ctx, authz, repo, bookmark, additional_changesets)
            .await?;

        Ok(())
    }

    async fn check_extras(
        &mut self,
        ctx: &CoreContext,
        repo: &impl Repo,
        bookmark: &BookmarkKey,
        kind: BookmarkKind,
        additional_changesets: AdditionalChangesets,
    ) -> Result<(), BookmarkMovementError> {
        if (kind == BookmarkKind::Publishing || kind == BookmarkKind::PullDefaultPublishing)
            && !repo
                .repo_config()
                .pushrebase
                .allow_change_xrepo_mapping_extra
        {
            self.load_additional_changesets(ctx, repo, bookmark, additional_changesets)
                .await
                .context("Failed to load additional affected changesets")?;

            for bcs in self.iter() {
                if bcs
                    .hg_extra()
                    .any(|(name, _)| name == CHANGE_XREPO_MAPPING_EXTRA)
                {
                    // This extra is used in backsyncer, and it changes how commit
                    // is rewritten from a large repo to a small repo. This is dangerous
                    // operation, and we don't allow landing a commit with this extra set.
                    return Err(anyhow!(
                        "Disallowed extra {} is set on {}.",
                        CHANGE_XREPO_MAPPING_EXTRA,
                        bcs.get_changeset_id()
                    )
                    .into());
                }
            }
        }

        Ok(())
    }

    /// If the push is to a public bookmark, and the casefolding check is
    /// enabled, check that no affected changeset has case conflicts.
    async fn check_case_conflicts(
        &mut self,
        ctx: &CoreContext,
        repo: &impl Repo,
        bookmark: &BookmarkKey,
        kind: BookmarkKind,
        additional_changesets: AdditionalChangesets,
    ) -> Result<(), BookmarkMovementError> {
        if (kind == BookmarkKind::Publishing || kind == BookmarkKind::PullDefaultPublishing)
            && repo.repo_config().pushrebase.flags.casefolding_check
        {
            self.load_additional_changesets(ctx, repo, bookmark, additional_changesets)
                .await
                .context("Failed to load additional affected changesets")?;

            stream::iter(self.iter().map(Ok))
                .try_for_each_concurrent(100, |bcs| async move {
                    let bcs_id = bcs.get_changeset_id();

                    let sk_mf = repo
                        .repo_derived_data()
                        .derive::<RootSkeletonManifestId>(ctx, bcs_id)
                        .await
                        .map_err(Error::from)?
                        .into_skeleton_manifest_id()
                        .load(ctx, repo.repo_blobstore())
                        .await
                        .map_err(Error::from)?;
                    if sk_mf.has_case_conflicts() {
                        // We only reject a commit if it introduces new case
                        // conflicts compared to its parents.
                        let parents = stream::iter(bcs.parents().map(|parent_bcs_id| async move {
                            repo.repo_derived_data()
                                .derive::<RootSkeletonManifestId>(ctx, parent_bcs_id)
                                .await
                                .map_err(Error::from)?
                                .into_skeleton_manifest_id()
                                .load(ctx, repo.repo_blobstore())
                                .await
                                .map_err(Error::from)
                        }))
                        .buffered(10)
                        .try_collect::<Vec<_>>()
                        .await?;

                        if let Some((path1, path2)) = sk_mf
                            .first_new_case_conflict(ctx, repo.repo_blobstore(), parents)
                            .await?
                        {
                            return Err(BookmarkMovementError::CaseConflict {
                                changeset_id: bcs_id,
                                path1,
                                path2,
                            });
                        }
                    }
                    Ok(())
                })
                .await?;
        }
        Ok(())
    }

    /// If this is a user-initiated update to a public bookmark, run the
    /// hooks against the affected changesets. Also run hooks if it is a
    /// service-initiated pushrebase but hooks will run with taking this
    /// into account.
    async fn check_hooks(
        &mut self,
        ctx: &CoreContext,
        authz: &AuthorizationContext,
        repo: &impl Repo,
        hook_manager: &HookManager,
        bookmark: &BookmarkKey,
        pushvars: Option<&HashMap<String, Bytes>>,
        reason: BookmarkUpdateReason,
        kind: BookmarkKind,
        additional_changesets: AdditionalChangesets,
        cross_repo_push_source: CrossRepoPushSource,
    ) -> Result<(), BookmarkMovementError> {
        if (kind == BookmarkKind::Publishing || kind == BookmarkKind::PullDefaultPublishing)
            && should_run_hooks(authz, reason)
        {
            if reason == BookmarkUpdateReason::Push
                && tunables().disable_hooks_on_plain_push().unwrap_or_default()
            {
                // Skip running hooks for this plain push.
                return Ok(());
            }

            if hook_manager.hooks_exist_for_bookmark(bookmark) {
                self.load_additional_changesets(ctx, repo, bookmark, additional_changesets)
                    .await
                    .context("Failed to load additional affected changesets")?;

                let skip_running_hooks_if_public: bool = repo
                    .repo_bookmark_attrs()
                    .select(bookmark)
                    .map(|attr| attr.params().allow_move_to_public_commits_without_hooks)
                    .any(|x| x);
                if skip_running_hooks_if_public && !self.adding_new_changesets_to_repo() {
                    // For some bookmarks we allow to skip running hooks if:
                    // 1) this is just a bookmark move i.e. no new commits are added or pushrebased to the repo
                    // 2) we are allowed to skip commits for a bookmark like that
                    // 3) if all commits that are affectd by this bookmark move are public (which means
                    //  we should have already ran hooks for these commits).

                    let cs_ids = self
                        .iter()
                        .map(|bcs| bcs.get_changeset_id())
                        .collect::<Vec<_>>();

                    let public = repo
                        .phases()
                        .get_public(ctx, cs_ids.clone(), false /* ephemeral_derive */)
                        .await?;
                    if public == cs_ids.into_iter().collect::<HashSet<_>>() {
                        return Ok(());
                    }
                }

                if !self.is_empty() {
                    let push_authored_by = if authz.is_service() {
                        PushAuthoredBy::Service
                    } else {
                        PushAuthoredBy::User
                    };
                    run_hooks(
                        ctx,
                        hook_manager,
                        bookmark,
                        self.iter(),
                        pushvars,
                        cross_repo_push_source,
                        push_authored_by,
                    )
                    .await?;
                }
            }
        }
        Ok(())
    }

    /// Check whether the user has permissions to modify the paths that are
    /// modified by the changesets that are affected by the bookmark move.
    async fn check_path_permissions(
        &mut self,
        ctx: &CoreContext,
        authz: &AuthorizationContext,
        repo: &impl Repo,
        bookmark: &BookmarkKey,
        additional_changesets: AdditionalChangesets,
    ) -> Result<(), BookmarkMovementError> {
        // For optimization, first check if the user is permitted to modify
        // all paths.  In that case we don't need to find out which paths were
        // affected.
        if authz.check_any_path_write(ctx, repo).await.is_denied() {
            // User is not permitted to write to all paths, check if the paths
            // touched by the changesets are permitted.
            self.load_additional_changesets(ctx, repo, bookmark, additional_changesets)
                .await
                .context("Failed to load additional affected changesets")?;

            for cs in self.iter() {
                authz.require_changeset_paths_write(ctx, repo, cs).await?;
            }
        }
        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::HashMap;

use bookmarks::BookmarkUpdateReason;
use bookmarks_types::BookmarkKey;
use bookmarks_types::BookmarkKind;
use bytes::Bytes;
use context::CoreContext;
use hooks::CrossRepoPushSource;
use hooks::HookManager;
use mononoke_types::BonsaiChangeset;
use mononoke_types::ChangesetId;
use repo_authorization::AuthorizationContext;
use repo_authorization::RepoWriteOperation;
use repo_update_logger::find_draft_ancestors;
use repo_update_logger::log_bookmark_operation;
use repo_update_logger::log_new_bonsai_changesets;
use repo_update_logger::BookmarkInfo;
use repo_update_logger::BookmarkOperation;

use crate::affected_changesets::AdditionalChangesets;
use crate::affected_changesets::AffectedChangesets;
use crate::repo_lock::check_repo_lock;
use crate::restrictions::check_bookmark_sync_config;
use crate::restrictions::BookmarkKindRestrictions;
use crate::BookmarkMovementError;
use crate::Repo;

#[must_use = "CreateBookmarkOp must be run to have an effect"]
pub struct CreateBookmarkOp<'op> {
    bookmark: &'op BookmarkKey,
    target: ChangesetId,
    reason: BookmarkUpdateReason,
    kind_restrictions: BookmarkKindRestrictions,
    cross_repo_push_source: CrossRepoPushSource,
    affected_changesets: AffectedChangesets,
    pushvars: Option<&'op HashMap<String, Bytes>>,
    log_new_public_commits_to_scribe: bool,
    only_log_acl_checks: bool,
}

impl<'op> CreateBookmarkOp<'op> {
    pub fn new(
        bookmark: &'op BookmarkKey,
        target: ChangesetId,
        reason: BookmarkUpdateReason,
    ) -> CreateBookmarkOp<'op> {
        CreateBookmarkOp {
            bookmark,
            target,
            reason,
            kind_restrictions: BookmarkKindRestrictions::AnyKind,
            cross_repo_push_source: CrossRepoPushSource::NativeToThisRepo,
            affected_changesets: AffectedChangesets::new(),
            pushvars: None,
            log_new_public_commits_to_scribe: false,
            only_log_acl_checks: false,
        }
    }

    pub fn only_if_scratch(mut self) -> Self {
        self.kind_restrictions = BookmarkKindRestrictions::OnlyScratch;
        self
    }

    pub fn only_if_public(mut self) -> Self {
        self.kind_restrictions = BookmarkKindRestrictions::OnlyPublishing;
        self
    }

    pub fn with_pushvars(mut self, pushvars: Option<&'op HashMap<String, Bytes>>) -> Self {
        self.pushvars = pushvars;
        self
    }

    pub fn log_new_public_commits_to_scribe(mut self) -> Self {
        self.log_new_public_commits_to_scribe = true;
        self
    }

    pub fn only_log_acl_checks(mut self, only_log: bool) -> Self {
        self.only_log_acl_checks = only_log;
        self
    }

    /// Include bonsai changesets for changesets that have just been added to
    /// the repository.
    pub fn with_new_changesets(
        mut self,
        changesets: HashMap<ChangesetId, BonsaiChangeset>,
    ) -> Self {
        self.affected_changesets.add_new_changesets(changesets);
        self
    }

    pub fn with_push_source(mut self, cross_repo_push_source: CrossRepoPushSource) -> Self {
        self.cross_repo_push_source = cross_repo_push_source;
        self
    }

    pub async fn run(
        mut self,
        ctx: &'op CoreContext,
        authz: &'op AuthorizationContext,
        repo: &'op impl Repo,
        hook_manager: &'op HookManager,
    ) -> Result<(), BookmarkMovementError> {
        let kind = self.kind_restrictions.check_kind(repo, self.bookmark)?;

        if self.only_log_acl_checks {
            if authz
                .check_repo_write(ctx, repo, RepoWriteOperation::CreateBookmark(kind))
                .await
                .is_denied()
            {
                ctx.scuba()
                    .clone()
                    .log_with_msg("Repo write ACL check would fail for bookmark create", None);
            }
        } else {
            authz
                .require_repo_write(ctx, repo, RepoWriteOperation::CreateBookmark(kind))
                .await?;
        }
        authz
            .require_bookmark_modify(ctx, repo, self.bookmark)
            .await?;

        check_bookmark_sync_config(repo, self.bookmark, kind)?;

        self.affected_changesets
            .check_restrictions(
                ctx,
                authz,
                repo,
                hook_manager,
                self.bookmark,
                self.pushvars,
                self.reason,
                kind,
                AdditionalChangesets::Ancestors(self.target),
                self.cross_repo_push_source,
            )
            .await?;

        check_repo_lock(repo, kind, self.pushvars, ctx.metadata().identities()).await?;

        let mut txn = repo.bookmarks().create_transaction(ctx.clone());
        let txn_hook;

        let commits_to_log = match kind {
            BookmarkKind::Scratch => {
                txn_hook = None;

                ctx.scuba()
                    .clone()
                    .add("bookmark", self.bookmark.to_string())
                    .log_with_msg("Creating scratch bookmark", None);

                txn.create_scratch(self.bookmark, self.target)?;
                vec![]
            }
            BookmarkKind::Publishing | BookmarkKind::PullDefaultPublishing => {
                crate::restrictions::check_restriction_ensure_ancestor_of(
                    ctx,
                    repo,
                    self.bookmark,
                    self.target,
                )
                .await?;

                let txn_hook_fut = crate::git_mapping::populate_git_mapping_txn_hook(
                    ctx,
                    repo,
                    self.target,
                    self.affected_changesets.new_changesets(),
                );

                let to_log = async {
                    if self.log_new_public_commits_to_scribe {
                        let res = find_draft_ancestors(ctx, repo, self.target).await;
                        match res {
                            Ok(bcss) => bcss,
                            Err(err) => {
                                ctx.scuba().clone().log_with_msg(
                                    "Failed to find draft ancestors",
                                    Some(format!("{}", err)),
                                );
                                vec![]
                            }
                        }
                    } else {
                        vec![]
                    }
                };

                let (txn_hook_res, to_log) = futures::join!(txn_hook_fut, to_log);
                txn_hook = txn_hook_res?;

                ctx.scuba()
                    .clone()
                    .add("bookmark", self.bookmark.to_string())
                    .log_with_msg("Creating public bookmark", None);

                txn.create(self.bookmark, self.target, self.reason)?;
                to_log
            }
        };

        let ok = match txn_hook {
            Some(txn_hook) => txn.commit_with_hook(txn_hook).await?,
            None => txn.commit().await?,
        };
        if !ok {
            return Err(BookmarkMovementError::TransactionFailed);
        }

        if self.log_new_public_commits_to_scribe {
            log_new_bonsai_changesets(ctx, repo, self.bookmark, kind, commits_to_log).await;
        }

        let info = BookmarkInfo {
            bookmark_name: self.bookmark.clone(),
            bookmark_kind: kind,
            operation: BookmarkOperation::Create(self.target),
            reason: self.reason,
        };
        log_bookmark_operation(ctx, repo, &info).await;
        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::HashMap;

use bookmarks::BookmarkUpdateReason;
use bookmarks_types::BookmarkKey;
use bookmarks_types::BookmarkKind;
use bytes::Bytes;
use context::CoreContext;
use mononoke_types::ChangesetId;
use repo_authorization::AuthorizationContext;
use repo_authorization::RepoWriteOperation;
use repo_update_logger::log_bookmark_operation;
use repo_update_logger::BookmarkInfo;
use repo_update_logger::BookmarkOperation;

use crate::repo_lock::check_repo_lock;
use crate::restrictions::check_bookmark_sync_config;
use crate::restrictions::BookmarkKindRestrictions;
use crate::BookmarkMovementError;
use crate::Repo;

#[must_use = "DeleteBookmarkOp must be run to have an effect"]
pub struct DeleteBookmarkOp<'op> {
    bookmark: &'op BookmarkKey,
    old_target: ChangesetId,
    reason: BookmarkUpdateReason,
    kind_restrictions: BookmarkKindRestrictions,
    pushvars: Option<&'op HashMap<String, Bytes>>,
    only_log_acl_checks: bool,
}

impl<'op> DeleteBookmarkOp<'op> {
    pub fn new(
        bookmark: &'op BookmarkKey,
        old_target: ChangesetId,
        reason: BookmarkUpdateReason,
    ) -> DeleteBookmarkOp<'op> {
        DeleteBookmarkOp {
            bookmark,
            old_target,
            reason,
            kind_restrictions: BookmarkKindRestrictions::AnyKind,
            pushvars: None,
            only_log_acl_checks: false,
        }
    }

    pub fn only_if_scratch(mut self) -> Self {
        self.kind_restrictions = BookmarkKindRestrictions::OnlyScratch;
        self
    }

    pub fn only_if_public(mut self) -> Self {
        self.kind_restrictions = BookmarkKindRestrictions::OnlyPublishing;
        self
    }

    pub fn with_pushvars(mut self, pushvars: Option<&'op HashMap<String, Bytes>>) -> Self {
        self.pushvars = pushvars;
        self
    }

    pub fn only_log_acl_checks(mut self, only_log: bool) -> Self {
        self.only_log_acl_checks = only_log;
        self
    }

    pub async fn run(
        self,
        ctx: &'op CoreContext,
        authz: &'op AuthorizationContext,
        repo: &'op impl Repo,
    ) -> Result<(), BookmarkMovementError> {
        let kind = self.kind_restrictions.check_kind(repo, self.bookmark)?;

        if self.only_log_acl_checks {
            if authz
                .check_repo_write(ctx, repo, RepoWriteOperation::DeleteBookmark(kind))
                .await
                .is_denied()
            {
                ctx.scuba()
                    .clone()
                    .log_with_msg("Repo write ACL check would fail for bookmark delete", None);
            }
        } else {
            authz
                .require_repo_write(ctx, repo, RepoWriteOperation::DeleteBookmark(kind))
                .await?;
        }
        authz
            .require_bookmark_modify(ctx, repo, self.bookmark)
            .await?;

        check_bookmark_sync_config(repo, self.bookmark, kind)?;

        if repo
            .repo_bookmark_attrs()
            .is_fast_forward_only(self.bookmark)
        {
            // Cannot delete fast-forward-only bookmarks.
            return Err(BookmarkMovementError::DeletionProhibited {
                bookmark: self.bookmark.clone(),
            });
        }

        check_repo_lock(repo, kind, self.pushvars, ctx.metadata().identities()).await?;

        ctx.scuba()
            .clone()
            .add("bookmark", self.bookmark.to_string())
            .log_with_msg("Deleting bookmark", None);
        let mut txn = repo.bookmarks().create_transaction(ctx.clone());
        match kind {
            BookmarkKind::Scratch => {
                txn.delete_scratch(self.bookmark, self.old_target)?;
            }
            BookmarkKind::Publishing | BookmarkKind::PullDefaultPublishing => {
                txn.delete(self.bookmark, self.old_target, self.reason)?;
            }
        }

        let ok = txn.commit().await?;
        if !ok {
            return Err(BookmarkMovementError::TransactionFailed);
        }

        let info = BookmarkInfo {
            bookmark_name: self.bookmark.clone(),
            bookmark_kind: kind,
            operation: BookmarkOperation::Delete(self.old_target),
            reason: self.reason,
        };
        log_bookmark_operation(ctx, repo, &info).await;
        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.
 */

//! Functions for maintaining git mappings during bookmark movement.

use std::borrow::Cow;
use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;

use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use blobstore::Loadable;
use bonsai_git_mapping::extract_git_sha1_from_bonsai_extra;
use bonsai_git_mapping::BonsaiGitMapping;
use bonsai_git_mapping::BonsaiGitMappingEntry;
use bookmarks::BookmarkTransactionHook;
use cloned::cloned;
use context::CoreContext;
use futures::future::try_join;
use futures::future::FutureExt;
use futures::future::TryFutureExt;
use futures::stream::FuturesOrdered;
use futures::stream::StreamExt;
use mononoke_types::BonsaiChangeset;
use mononoke_types::ChangesetId;

use crate::BookmarkMovementError;
use crate::Repo;

/// New mapping entries that should be added to the mapping.
struct NewMappingEntries {
    /// New entries to add.
    new_mapping_entries: Vec<BonsaiGitMappingEntry>,

    /// Count of entries that were created from newly uploaded changesets.
    from_new_changesets: usize,

    /// Count of entries that were created from ancestors with no mapping.
    from_ancestors_no_mapping: usize,
}

/// Description of a changeset for which a mapping needs to be generated.
struct ChangesetMappingNeeded<'a> {
    /// Whether knowledge of this need came from this changeset being
    /// a new changeset.
    from_new_changeset: bool,

    /// The changeset id that needs a mapping.
    cs_id: ChangesetId,

    /// The bonsai changeset that needs a mapping.
    bcs: Cow<'a, BonsaiChangeset>,
}

/// Find ancestors of `start` which have git mapping extras but do not
/// have a git mapping entry set in db, and generate their mapping entries.
async fn new_mapping_entries(
    ctx: &CoreContext,
    repo: &impl Repo,
    start: ChangesetId,
    new_changesets: &HashMap<ChangesetId, BonsaiChangeset>,
) -> Result<NewMappingEntries> {
    let mut new_mapping_entries = Vec::new();
    let mut from_new_changesets = 0;
    let mut from_ancestors_no_mapping = 0;

    let mut visited = HashSet::new();
    let mut queue = FuturesOrdered::new();
    let mut get_new_queue_entry = |cs_id: ChangesetId| {
        if visited.insert(cs_id) {
            Some(async move {
                if let Some(bcs) = new_changesets.get(&cs_id) {
                    Ok::<_, Error>(Some(ChangesetMappingNeeded {
                        from_new_changeset: true,
                        cs_id,
                        bcs: Cow::Borrowed(bcs),
                    }))
                } else {
                    let bcs_fut = {
                        cloned!(ctx, repo);
                        async move {
                            cs_id
                                .load(&ctx, repo.repo_blobstore())
                                .map_err(Error::from)
                                .await
                        }
                    };

                    let mapping_fut = repo.bonsai_git_mapping().get(ctx, cs_id.into());

                    let (bcs, git_mapping) = try_join(bcs_fut, mapping_fut).await?;
                    if git_mapping.is_empty() {
                        Ok(Some(ChangesetMappingNeeded {
                            from_new_changeset: false,
                            cs_id,
                            bcs: Cow::Owned(bcs),
                        }))
                    } else {
                        // The mapping is already known for this changeset, so
                        // generating the mapping is not needed.
                        Ok(None)
                    }
                }
            })
        } else {
            None
        }
    };

    queue.extend(get_new_queue_entry(start));

    while let Some(entry) = queue.next().await {
        let needed = match entry? {
            Some(cs_mapping_needed) => cs_mapping_needed,
            None => {
                // The mapping is already known for this entry, so we can stop
                // traversing here.
                continue;
            }
        };

        let git_sha1 = match extract_git_sha1_from_bonsai_extra(needed.bcs.hg_extra())
            .with_context(|| format!("Failed to extract Git Sha1 from {}", needed.cs_id))?
        {
            Some(git_sha1) => git_sha1,
            None => {
                // Don't traverse past commits that do not have git sha1 set
                // This is done deliberately to avoid retraversing these commits over
                // and over each time new mappings are added.
                continue;
            }
        };

        for p in needed.bcs.parents() {
            queue.extend(get_new_queue_entry(p));
        }
        if needed.from_new_changeset {
            from_new_changesets += 1;
        } else {
            from_ancestors_no_mapping += 1;
        }
        let mapping_entry = BonsaiGitMappingEntry {
            git_sha1,
            bcs_id: needed.cs_id,
        };
        new_mapping_entries.push(mapping_entry);
    }

    Ok(NewMappingEntries {
        new_mapping_entries,
        from_new_changesets,
        from_ancestors_no_mapping,
    })
}

fn upload_mapping_entries_bookmark_txn_hook(
    bonsai_git_mapping: Arc<dyn BonsaiGitMapping>,
    entries: NewMappingEntries,
) -> BookmarkTransactionHook {
    let entries = Arc::new(entries);
    Arc::new(move |ctx, sql_txn| {
        // We expect new_changesets + ancestors_no_mapping = inserting.  For
        // pushes, ancestors_no_mapping should be zero.  For requests made via
        // APIs, it may be non-zero.
        ctx.scuba()
            .clone()
            .add("git_mapping_inserting", entries.new_mapping_entries.len())
            .add("git_mapping_new_changesets", entries.from_new_changesets)
            .add(
                "git_mapping_ancestors_no_mapping",
                entries.from_ancestors_no_mapping,
            )
            .log_with_msg("Inserting git mapping", None);

        let bonsai_git_mapping = bonsai_git_mapping.clone();
        let mapping_entries = entries.new_mapping_entries.clone();
        async move {
            let sql_txn = bonsai_git_mapping
                .bulk_add_git_mapping_in_transaction(&ctx, mapping_entries.as_slice(), sql_txn)
                .map_err(Error::from)
                .await?;
            ctx.scuba()
                .clone()
                .log_with_msg("Inserted git mapping", None);
            Ok(sql_txn)
        }
        .boxed()
    })
}

/// Generate a bookmark transaction hook that will populate the git mapping
/// with new entries for the new mapped commits reachable from `new_head`.
///
/// Commits that have just been added to the repository (and so their
/// bonsai changeset is known and it is also known that they do not have a
/// mapping) can be passed in via `new_changesets` to skip checking the
/// mapping for whether the changeset is already present and to avoid
/// fetching the bonsai changeset.
pub(crate) async fn populate_git_mapping_txn_hook(
    ctx: &CoreContext,
    repo: &impl Repo,
    new_head: ChangesetId,
    new_changesets: &HashMap<ChangesetId, BonsaiChangeset>,
) -> Result<Option<BookmarkTransactionHook>, BookmarkMovementError> {
    if repo.repo_config().pushrebase.populate_git_mapping {
        let entries = new_mapping_entries(ctx, repo, new_head, new_changesets).await?;
        Ok(Some(upload_mapping_entries_bookmark_txn_hook(
            repo.bonsai_git_mapping_arc(),
            entries,
        )))
    } else {
        Ok(None)
    }
}

#[cfg(test)]
mod tests {
    use anyhow::Result;
    use blobrepo::AsBlobRepo;
    use blobrepo::BlobRepo;
    use bonsai_git_mapping::CONVERT_REVISION_EXTRA;
    use bonsai_git_mapping::HGGIT_SOURCE_EXTRA;
    use bookmarks::BookmarkKey;
    use bookmarks::BookmarkUpdateReason;
    use bookmarks::BookmarksRef;
    use borrowed::borrowed;
    use fbinit::FacebookInit;
    use maplit::hashmap;
    use maplit::hashset;
    use mononoke_api_types::InnerRepo;
    use mononoke_types::hash::GitSha1;
    use mononoke_types_mocks::hash::FIVES_GIT_SHA1;
    use mononoke_types_mocks::hash::FOURS_GIT_SHA1;
    use mononoke_types_mocks::hash::ONES_GIT_SHA1;
    use mononoke_types_mocks::hash::SIXES_GIT_SHA1;
    use mononoke_types_mocks::hash::THREES_GIT_SHA1;
    use mononoke_types_mocks::hash::TWOS_GIT_SHA1;
    use repo_blobstore::RepoBlobstoreRef;
    use test_repo_factory::TestRepoFactory;
    use tests_utils::drawdag::changes;
    use tests_utils::drawdag::create_from_dag_with_changes;
    use tests_utils::CreateCommitContext;

    use super::*;

    fn add_git_extras(
        context: CreateCommitContext<BlobRepo>,
        hash: GitSha1,
    ) -> CreateCommitContext<BlobRepo> {
        context
            .add_extra(
                CONVERT_REVISION_EXTRA.to_string(),
                format!("{}", hash).as_bytes().to_vec(),
            )
            .add_extra(HGGIT_SOURCE_EXTRA.to_string(), b"git".to_vec())
    }

    fn mapping_entries(entries: &[BonsaiGitMappingEntry]) -> HashSet<(ChangesetId, GitSha1)> {
        entries
            .iter()
            .map(|entry| (entry.bcs_id, entry.git_sha1))
            .collect()
    }

    async fn apply_entries(
        ctx: &CoreContext,
        repo: &impl Repo,
        bookmark: &BookmarkKey,
        old_target: ChangesetId,
        new_target: ChangesetId,
        entries: NewMappingEntries,
    ) -> Result<()> {
        let mut txn = repo.bookmarks().create_transaction(ctx.clone());
        txn.update(
            bookmark,
            new_target,
            old_target,
            BookmarkUpdateReason::TestMove,
        )?;
        let ok = txn
            .commit_with_hook(upload_mapping_entries_bookmark_txn_hook(
                repo.bonsai_git_mapping_arc(),
                entries,
            ))
            .await?;
        assert!(ok);
        Ok(())
    }

    #[fbinit::test]
    async fn test_new_mapping_entries(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let repo: InnerRepo = TestRepoFactory::new(fb)?.build().await?;
        let bookmark = BookmarkKey::new("main")?;
        borrowed!(ctx, repo);

        let dag = create_from_dag_with_changes(
            ctx,
            repo.as_blob_repo(),
            r##"
                Z-A-B-C-D
                     \
                      X-Y
                       \
                        E-F
            "##,
            changes! {
                "A" => |c| add_git_extras(c, ONES_GIT_SHA1),
                "B" => |c| add_git_extras(c, TWOS_GIT_SHA1),
                "C" => |c| add_git_extras(c, THREES_GIT_SHA1),
                "D" => |c| add_git_extras(c, FOURS_GIT_SHA1),
                "E" => |c| add_git_extras(c, FIVES_GIT_SHA1),
                "F" => |c| add_git_extras(c, SIXES_GIT_SHA1),
            },
        )
        .await?;

        let a = *dag.get("A").unwrap();
        let b = *dag.get("B").unwrap();
        let c = *dag.get("C").unwrap();
        let d = *dag.get("D").unwrap();
        let e = *dag.get("E").unwrap();
        let f = *dag.get("F").unwrap();
        let y = *dag.get("Y").unwrap();
        let z = *dag.get("Z").unwrap();
        let a_bcs = a.load(ctx, repo.repo_blobstore()).await?;
        let b_bcs = b.load(ctx, repo.repo_blobstore()).await?;
        let f_bcs = f.load(ctx, repo.repo_blobstore()).await?;
        let y_bcs = y.load(ctx, repo.repo_blobstore()).await?;

        let mut txn = repo.bookmarks().create_transaction(ctx.clone());
        txn.create(&bookmark, z, BookmarkUpdateReason::TestMove)?;
        let ok = txn.commit().await?;
        assert!(ok);

        // Initial create with two changesets.
        let entries = new_mapping_entries(
            ctx,
            repo,
            b,
            &hashmap! {
                a => a_bcs,
                b => b_bcs,
            },
        )
        .await?;

        assert_eq!(
            mapping_entries(&entries.new_mapping_entries),
            hashset! { (a, ONES_GIT_SHA1), (b, TWOS_GIT_SHA1) },
        );
        assert_eq!(entries.from_new_changesets, 2);
        assert_eq!(entries.from_ancestors_no_mapping, 0);

        apply_entries(ctx, repo, &bookmark, z, b, entries).await?;

        // Addition using existing changesets.
        let entries = new_mapping_entries(ctx, repo, d, &hashmap! {}).await?;
        assert_eq!(
            mapping_entries(&entries.new_mapping_entries),
            hashset! { (c, THREES_GIT_SHA1), (d, FOURS_GIT_SHA1) },
        );
        assert_eq!(entries.from_new_changesets, 0);
        assert_eq!(entries.from_ancestors_no_mapping, 2);

        apply_entries(ctx, repo, &bookmark, b, d, entries).await?;

        // Move to commits with no mapping.
        let entries = new_mapping_entries(ctx, repo, y, &hashmap! {y => y_bcs}).await?;
        assert_eq!(mapping_entries(&entries.new_mapping_entries), hashset! {});
        assert_eq!(entries.from_new_changesets, 0);
        assert_eq!(entries.from_ancestors_no_mapping, 0);

        apply_entries(ctx, repo, &bookmark, d, y, entries).await?;

        // Move to descendants of commit with no mapping.
        let entries = new_mapping_entries(ctx, repo, f, &hashmap! {f => f_bcs}).await?;
        assert_eq!(
            mapping_entries(&entries.new_mapping_entries),
            hashset! {
                (e, FIVES_GIT_SHA1), (f, SIXES_GIT_SHA1)
            }
        );
        assert_eq!(entries.from_new_changesets, 1);
        assert_eq!(entries.from_ancestors_no_mapping, 1);

        apply_entries(ctx, repo, &bookmark, y, f, entries).await?;

        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::HashMap;

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Result;
use bookmarks_types::BookmarkKey;
use bytes::Bytes;
use context::CoreContext;
use futures_stats::TimedFutureExt;
use hooks::CrossRepoPushSource;
use hooks::HookManager;
use hooks::HookOutcome;
use hooks::PushAuthoredBy;
use mononoke_types::BonsaiChangeset;
use tunables::tunables;

use crate::BookmarkMovementError;

pub async fn is_admin_bypass(
    ctx: &CoreContext,
    hook_manager: &HookManager,
    pushvars: Option<&HashMap<String, Bytes>>,
) -> Result<bool> {
    let pushvars = match pushvars {
        Some(pushvars) => pushvars,
        None => {
            return Ok(false);
        }
    };

    if !pushvars.contains_key("BYPASS_ALL_HOOKS") {
        return Ok(false);
    }

    if !hook_manager
        .get_admin_perm_checker()
        .is_member(ctx.metadata().identities())
        .await
    {
        return Err(anyhow!(
            "In order to use BYPASS_ALL_HOOKS pushvar one needs to be member of the scm group."
        ));
    }

    Ok(true)
}

fn take_n_changeset_ids<'a>(
    changesets: impl Iterator<Item = &'a BonsaiChangeset> + Clone,
    n: usize,
) -> Vec<String> {
    changesets
        .map(|cs| cs.get_changeset_id().to_string())
        .take(n)
        .collect()
}

pub async fn run_hooks(
    ctx: &CoreContext,
    hook_manager: &HookManager,
    bookmark: &BookmarkKey,
    changesets: impl Iterator<Item = &BonsaiChangeset> + Clone,
    pushvars: Option<&HashMap<String, Bytes>>,
    cross_repo_push_source: CrossRepoPushSource,
    push_authored_by: PushAuthoredBy,
) -> Result<(), BookmarkMovementError> {
    if cross_repo_push_source == CrossRepoPushSource::PushRedirected {
        if tunables()
            .disable_running_hooks_in_pushredirected_repo()
            .unwrap_or_default()
        {
            let cs_ids = take_n_changeset_ids(changesets, 10);
            ctx.scuba()
                .clone()
                .add("bookmark", bookmark.to_string())
                .add("changesets", cs_ids)
                .log_with_msg("Hook execution in pushredirected repo was disabled", None);
            return Ok(());
        }
    }

    if is_admin_bypass(ctx, hook_manager, pushvars).await? || hook_manager.all_hooks_bypassed() {
        let mut scuba_bypassed_commits = hook_manager.scuba_bypassed_commits().clone();
        let cs_ids = take_n_changeset_ids(changesets, 10);

        scuba_bypassed_commits
            .add_metadata(ctx.metadata())
            .add("bookmark", bookmark.to_string())
            .add("changesets", cs_ids)
            .add("repo_name", hook_manager.repo_name().clone());

        if let Some(pushvars) = pushvars {
            scuba_bypassed_commits.add(
                "pushvars",
                pushvars
                    .iter()
                    .map(|(key, val)| format!("{}={:?}", key, val))
                    .collect::<Vec<_>>(),
            );
        }

        scuba_bypassed_commits
            .log_with_msg("Bypassed all hooks using BYPASS_ALL_HOOKS pushvar.", None);
        return Ok(());
    }

    let (stats, outcomes) = hook_manager
        .run_hooks_for_bookmark(
            ctx,
            changesets,
            bookmark,
            pushvars,
            cross_repo_push_source,
            push_authored_by,
        )
        .timed()
        .await;
    let outcomes = outcomes.with_context(|| format!("Failed to run hooks for {}", bookmark))?;

    let rejections: Vec<_> = outcomes
        .into_iter()
        .filter_map(HookOutcome::into_rejection)
        .collect();

    ctx.scuba()
        .clone()
        .add_future_stats(&stats)
        .add("hook_rejections", rejections.len())
        .log_with_msg("Executed hooks", None);

    if rejections.is_empty() {
        Ok(())
    } else {
        Err(BookmarkMovementError::HookFailure(rejections))
    }
}